From fbd1977152c8fb133de341cc33fc5f027bd393b6 Mon Sep 17 00:00:00 2001
From: ronger <ronger-x@outlook.com>
Date: Fri, 26 May 2023 18:42:36 +0800
Subject: [PATCH] =?UTF-8?q?:art:=20=E4=BD=BF=E7=94=A8=E5=AE=A2=E6=88=B7?=
 =?UTF-8?q?=E7=AB=AF=20IP=20=E8=BF=9B=E8=A1=8C=E8=AF=B7=E6=B1=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../openai/service/IpAddressInterceptor.java  | 29 +++++++++++++++++++
 .../forest/openai/service/OpenAiService.java  | 10 +++++--
 2 files changed, 37 insertions(+), 2 deletions(-)
 create mode 100644 src/main/java/com/rymcu/forest/openai/service/IpAddressInterceptor.java

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();