diff --git a/src/main/java/com/rymcu/forest/openai/service/IpAddressInterceptor.java b/src/main/java/com/rymcu/forest/openai/service/IpAddressInterceptor.java new file mode 100644 index 0000000..ea8aab7 --- /dev/null +++ b/src/main/java/com/rymcu/forest/openai/service/IpAddressInterceptor.java @@ -0,0 +1,29 @@ +package com.rymcu.forest.openai.service; + +import okhttp3.Interceptor; +import okhttp3.Request; +import okhttp3.Response; + +import java.io.IOException; + +/** + * OkHttp Interceptor that adds an ip address header + * @author ronger + */ +public class IpAddressInterceptor implements Interceptor { + + private final String ip; + + IpAddressInterceptor(String ip) { + this.ip = ip; + } + + @Override + public Response intercept(Chain chain) throws IOException { + Request request = chain.request() + .newBuilder() + .header("x-forwarded-for", ip) + .build(); + return chain.proceed(request); + } +} diff --git a/src/main/java/com/rymcu/forest/openai/service/OpenAiService.java b/src/main/java/com/rymcu/forest/openai/service/OpenAiService.java index ad468b8..b9d00fb 100644 --- a/src/main/java/com/rymcu/forest/openai/service/OpenAiService.java +++ b/src/main/java/com/rymcu/forest/openai/service/OpenAiService.java @@ -5,6 +5,7 @@ import com.fasterxml.jackson.databind.DeserializationFeature; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.PropertyNamingStrategy; import com.rymcu.forest.util.SpringContextHolder; +import com.rymcu.forest.util.Utils; import com.theokanning.openai.DeleteResult; import com.theokanning.openai.OpenAiApi; import com.theokanning.openai.OpenAiError; @@ -30,18 +31,20 @@ import com.theokanning.openai.image.ImageResult; import com.theokanning.openai.model.Model; import com.theokanning.openai.moderation.ModerationRequest; import com.theokanning.openai.moderation.ModerationResult; - import io.reactivex.BackpressureStrategy; import io.reactivex.Flowable; import io.reactivex.Single; import okhttp3.*; import org.springframework.core.env.Environment; +import org.springframework.web.context.request.RequestContextHolder; +import org.springframework.web.context.request.ServletRequestAttributes; +import retrofit2.Call; import retrofit2.HttpException; import retrofit2.Retrofit; -import retrofit2.Call; import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory; import retrofit2.converter.jackson.JacksonConverterFactory; +import javax.servlet.http.HttpServletRequest; import java.io.IOException; import java.time.Duration; import java.util.List; @@ -336,8 +339,11 @@ public class OpenAiService { } public static OkHttpClient defaultClient(String token, Duration timeout) { + HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest(); + String ip = Utils.getIpAddress(request); return new OkHttpClient.Builder() .addInterceptor(new AuthenticationInterceptor(token)) + .addInterceptor(new IpAddressInterceptor(ip)) .connectionPool(new ConnectionPool(5, 1, TimeUnit.SECONDS)) .readTimeout(timeout.toMillis(), TimeUnit.MILLISECONDS) .build();