1 (aurora-develop#78)

* Add CORS middleware to handle cross-origin requests (aurora-develop#72) * Revert "Add CORS middleware to handle cross-origin requests (aurora-develop#72)" (aurora-develop#73) This reverts commit 03173fe. * Test au (aurora-develop#75) * Add CORS middleware to handle cross-origin requests * chroe: 抽离 baseurl * Update main.go * update: 更新 api 路径为/backend-anon * Refactor code to improve performance and readability * uuid (aurora-develop#74) * Add CORS middleware to handle cross-origin requests (aurora-develop#72) * Revert "Add CORS middleware to handle cross-origin requests (aurora-develop#72)" (aurora-develop#73) This reverts commit 03173fe. * 修复uuid失效，替换失效uuid --------- Co-authored-by: xiaofei <[email protected]> --------- Co-authored-by: aurorax-neo <[email protected]> Co-authored-by: learnLi <[email protected]> * 全局变量私有化 & 重排序导入包 (aurora-develop#76) * style: sort imports * refactor: 私有化全局变量 * Update build_docker.yml (aurora-develop#77) --------- Co-authored-by: aurorax-neo <[email protected]> Co-authored-by: learnLi <[email protected]> Co-authored-by: icpd <[email protected]>
okok110 · Apr 4, 2024 · 0b13f49 · 0b13f49
1 parent bf3d252
commit 0b13f49
Show file tree

Hide file tree

Showing 11 changed files with 131 additions and 104 deletions.
diff --git a/.github/workflows/build_docker.yml b/.github/workflows/build_docker.yml
@@ -33,7 +33,9 @@ jobs:
         uses: actions/cache@v2
         with:
           path: /tmp/.buildx-cache
-          key: ${{ runner.os }}-buildx-cache
+          key: ${{ runner.os }}-buildx-${{ github.sha }}
+          restore-keys: |
+            ${{ runner.os }}-buildx-
 
       - name: Build and push to GHCR
         uses: docker/build-push-action@v2
@@ -42,6 +44,8 @@ jobs:
           platforms: linux/amd64,linux/arm64
           file: Dockerfile
           push: true
-          tags: ${{ env.GHCR_REPO }}:latest
+          tags: |
+            ${{ env.GHCR_REPO }}:latest
+            ${{ env.GHCR_REPO }}:${{ github.sha }}
           cache-from: type=local,src=/tmp/.buildx-cache
-          cache-to: type=local,dest=/tmp/.buildx-cache,mode=max
+          cache-to: type=local,dest=/tmp/.buildx-cache,mode=max
diff --git a/auth.go b/auth.go
@@ -1,16 +1,17 @@
 package main
 
 import (
-	"aurora/internal/tokens"
 	"bufio"
 	"fmt"
 	"os"
 	"strconv"
 
+	"aurora/internal/tokens"
+
 	"github.com/google/uuid"
 )
 
-func readAccessToken() {
+func readAccessToken() *tokens.AccessToken {
 	var Secrets []*tokens.Secret
 	// Read accounts.txt and create a list of accounts
 	if _, err := os.Stat("access_tokens.txt"); err == nil {
@@ -62,5 +63,6 @@ func readAccessToken() {
 		}
 	}
 
-	ACCESS_TOKENS = tokens.NewAccessToken(Secrets)
+	token := tokens.NewAccessToken(Secrets)
+	return &token
 }
diff --git a/conversion/requests/chatgpt/convert.go b/conversion/requests/chatgpt/convert.go
@@ -1,12 +1,14 @@
 package chatgpt
 
 import (
+	"fmt"
+	"strings"
+
 	"aurora/internal/tokens"
 	chatgpt_types "aurora/typings/chatgpt"
 	official_types "aurora/typings/official"
-	"fmt"
+
 	arkose "github.com/xqdoo00o/funcaptcha"
-	"strings"
 )
 
 func ConvertAPIRequest(api_request official_types.APIRequest, secret *tokens.Secret, requireArk bool, proxy string) chatgpt_types.ChatGPTRequest {

diff --git a/conversion/response/chatgpt/convert.go b/conversion/response/chatgpt/convert.go
@@ -1,10 +1,11 @@
 package chatgpt
 
 import (
+	"strings"
+
 	"aurora/typings"
 	chatgpt_types "aurora/typings/chatgpt"
 	official_types "aurora/typings/official"
-	"strings"
 )
 
 func ConvertToString(chatgpt_response *chatgpt_types.ChatGPTResponse, previous_text *typings.StringStruct, role bool) string {

diff --git a/handlers.go b/handlers.go
@@ -1,17 +1,28 @@
 package main
 
 import (
-	chatgpt_request_converter "aurora/conversion/requests/chatgpt"
+	"strings"
+
+	chatgptrequestconverter "aurora/conversion/requests/chatgpt"
 	"aurora/httpclient"
 	"aurora/internal/chatgpt"
-
-	official_types "aurora/typings/official"
-	"strings"
+	"aurora/internal/proxys"
+	"aurora/internal/tokens"
+	officialtypes "aurora/typings/official"
 
 	"github.com/gin-gonic/gin"
 	"github.com/google/uuid"
 )
 
+type Handler struct {
+	proxy *proxys.IProxy
+	token *tokens.AccessToken
+}
+
+func NewHandle(proxy *proxys.IProxy, token *tokens.AccessToken) *Handler {
+	return &Handler{proxy: proxy, token: token}
+}
+
 func optionsHandler(c *gin.Context) {
 	// Set headers for CORS
 	c.Header("Access-Control-Allow-Origin", "*")
@@ -22,9 +33,9 @@ func optionsHandler(c *gin.Context) {
 	})
 }
 
-func refresh_handler(c *gin.Context) {
-	var refresh_token official_types.OpenAIRefreshToken
-	err := c.BindJSON(&refresh_token)
+func (h *Handler) refresh(c *gin.Context) {
+	var refreshToken officialtypes.OpenAIRefreshToken
+	err := c.BindJSON(&refreshToken)
 	if err != nil {
 		c.JSON(400, gin.H{"error": gin.H{
 			"message": "Request must be proper JSON",
@@ -34,8 +45,8 @@ func refresh_handler(c *gin.Context) {
 		}})
 		return
 	}
-	proxy_url := ProxyIP.GetProxyIP()
-	openaiRefreshToken, status, err := chatgpt.GETTokenForRefreshToken(refresh_token.RefreshToken, proxy_url)
+	proxyUrl := h.proxy.GetProxyIP()
+	openaiRefreshToken, status, err := chatgpt.GETTokenForRefreshToken(refreshToken.RefreshToken, proxyUrl)
 	if err != nil {
 		c.JSON(400, gin.H{
 			"message": "Request must be proper JSON",
@@ -48,9 +59,9 @@ func refresh_handler(c *gin.Context) {
 	c.JSON(status, openaiRefreshToken)
 }
 
-func session_handler(c *gin.Context) {
-	var session_token official_types.OpenAISessionToken
-	err := c.BindJSON(&session_token)
+func (h *Handler) session(c *gin.Context) {
+	var sessionToken officialtypes.OpenAISessionToken
+	err := c.BindJSON(&sessionToken)
 	if err != nil {
 		c.JSON(400, gin.H{
 			"message": "Request must be proper JSON",
@@ -60,8 +71,8 @@ func session_handler(c *gin.Context) {
 		})
 		return
 	}
-	proxy_url := ProxyIP.GetProxyIP()
-	openaiSessionToken, status, err := chatgpt.GETTokenForSessionToken(session_token.SessionToken, proxy_url)
+	proxy_url := h.proxy.GetProxyIP()
+	openaiSessionToken, status, err := chatgpt.GETTokenForSessionToken(sessionToken.SessionToken, proxy_url)
 	if err != nil {
 		c.JSON(400, gin.H{"error": gin.H{
 			"message": "Request must be proper JSON",
@@ -74,8 +85,8 @@ func session_handler(c *gin.Context) {
 	c.JSON(status, openaiSessionToken)
 }
 
-func nightmare(c *gin.Context) {
-	var original_request official_types.APIRequest
+func (h *Handler) nightmare(c *gin.Context) {
+	var original_request officialtypes.APIRequest
 	err := c.BindJSON(&original_request)
 	if err != nil {
 		c.JSON(400, gin.H{"error": gin.H{
@@ -87,13 +98,13 @@ func nightmare(c *gin.Context) {
 		return
 	}
 
-	proxy_url := ProxyIP.GetProxyIP()
-	secret := ACCESS_TOKENS.GetSecret()
+	proxyUrl := h.proxy.GetProxyIP()
+	secret := h.token.GetSecret()
 	authHeader := c.GetHeader("Authorization")
 	if authHeader != "" {
 		customAccessToken := strings.Replace(authHeader, "Bearer ", "", 1)
 		if strings.HasPrefix(customAccessToken, "eyJhbGciOiJSUzI1NiI") {
-			secret = ACCESS_TOKENS.GenerateTempToken(customAccessToken)
+			secret = h.token.GenerateTempToken(customAccessToken)
 		}
 	}
 	if secret == nil {
@@ -104,7 +115,7 @@ func nightmare(c *gin.Context) {
 
 	client := httpclient.NewStdClient()
 	uid := uuid.NewString()
-	turnStile, status, err := chatgpt.InitTurnStile(client, secret, proxy_url)
+	turnStile, status, err := chatgpt.InitTurnStile(client, secret, proxyUrl)
 	if err != nil {
 		c.JSON(status, gin.H{
 			"message": err.Error(),
@@ -115,17 +126,17 @@ func nightmare(c *gin.Context) {
 		return
 	}
 	if !secret.IsFree {
-		err = chatgpt.InitWSConn(secret.Token, uid, proxy_url)
+		err = chatgpt.InitWSConn(secret.Token, uid, proxyUrl)
 		if err != nil {
 			c.JSON(500, gin.H{"error": "unable to create ws tunnel"})
 			return
 		}
 	}
 
 	// Convert the chat request to a ChatGPT request
-	translated_request := chatgpt_request_converter.ConvertAPIRequest(original_request, secret, turnStile.Arkose, proxy_url)
+	translated_request := chatgptrequestconverter.ConvertAPIRequest(original_request, secret, turnStile.Arkose, proxyUrl)
 
-	response, err := chatgpt.POSTconversation(client, translated_request, secret, turnStile, proxy_url)
+	response, err := chatgpt.POSTconversation(client, translated_request, secret, turnStile, proxyUrl)
 	if err != nil {
 		c.JSON(500, gin.H{
 			"error": "error sending request",
@@ -151,9 +162,9 @@ func nightmare(c *gin.Context) {
 		translated_request.ParentMessageID = continue_info.ParentID
 
 		if turnStile.Arkose {
-			chatgpt_request_converter.RenewTokenForRequest(&translated_request, secret.PUID, proxy_url)
+			chatgptrequestconverter.RenewTokenForRequest(&translated_request, secret.PUID, proxyUrl)
 		}
-		response, err = chatgpt.POSTconversation(client, translated_request, secret, turnStile, proxy_url)
+		response, err = chatgpt.POSTconversation(client, translated_request, secret, turnStile, proxyUrl)
 		if err != nil {
 			c.JSON(500, gin.H{
 				"error": "error sending request",
@@ -169,30 +180,30 @@ func nightmare(c *gin.Context) {
 		return
 	}
 	if !original_request.Stream {
-		c.JSON(200, official_types.NewChatCompletion(full_response))
+		c.JSON(200, officialtypes.NewChatCompletion(full_response))
 	} else {
 		c.String(200, "data: [DONE]\n\n")
 	}
 	chatgpt.UnlockSpecConn(secret.Token, uid)
 }
 
-func engines_handler(c *gin.Context) {
-	proxy_url := ProxyIP.GetProxyIP()
-	secret := ACCESS_TOKENS.GetSecret()
+func (h *Handler) engines(c *gin.Context) {
+	proxyUrl := h.proxy.GetProxyIP()
+	secret := h.token.GetSecret()
 	authHeader := c.GetHeader("Authorization")
 	if authHeader != "" {
 		customAccessToken := strings.Replace(authHeader, "Bearer ", "", 1)
 		// Check if customAccessToken starts with sk-
 		if strings.HasPrefix(customAccessToken, "eyJhbGciOiJSUzI1NiI") {
-			secret = ACCESS_TOKENS.GenerateTempToken(customAccessToken)
+			secret = h.token.GenerateTempToken(customAccessToken)
 		}
 	}
-	if secret.Token == "" || secret == nil {
+	if secret == nil || secret.Token == "" {
 		c.JSON(400, gin.H{"error": "Not Account Found."})
 		return
 	}
 
-	resp, status, err := chatgpt.GETengines(secret, proxy_url)
+	resp, status, err := chatgpt.GETengines(secret, proxyUrl)
 	if err != nil {
 		c.JSON(500, gin.H{
 			"error": "error sending request",

diff --git a/httpclient/client.go b/httpclient/client.go
@@ -1,12 +1,14 @@
 package httpclient
 
 import (
-	"aurora/util"
 	"crypto/tls"
-	browser "github.com/EDDYCJY/fake-useragent"
-	"github.com/go-resty/resty/v2"
 	"net/http"
 	"time"
+
+	"aurora/util"
+
+	browser "github.com/EDDYCJY/fake-useragent"
+	"github.com/go-resty/resty/v2"
 )
 
 type RestyClient struct {

diff --git a/internal/chatgpt/request.go b/internal/chatgpt/request.go
@@ -1,10 +1,6 @@
 package chatgpt
 
 import (
-	"aurora/httpclient"
-	"aurora/internal/tokens"
-	"aurora/typings"
-	chatgpt_types "aurora/typings/chatgpt"
 	"bufio"
 	"bytes"
 	"context"
@@ -22,18 +18,19 @@ import (
 	"sync"
 	"time"
 
-	"github.com/aurorax-neo/go-logger"
-
-	"github.com/gorilla/websocket"
+	chatgpt_response_converter "aurora/conversion/response/chatgpt"
+	"aurora/httpclient"
+	"aurora/internal/tokens"
+	"aurora/typings"
+	chatgpt_types "aurora/typings/chatgpt"
+	official_types "aurora/typings/official"
 
+	"github.com/aurorax-neo/go-logger"
 	fhttp "github.com/bogdanfinn/fhttp"
 	tls_client "github.com/bogdanfinn/tls-client"
 	"github.com/bogdanfinn/tls-client/profiles"
 	"github.com/gin-gonic/gin"
-
-	chatgpt_response_converter "aurora/conversion/response/chatgpt"
-
-	official_types "aurora/typings/official"
+	"github.com/gorilla/websocket"
 )
 
 type connInfo struct {

diff --git a/internal/proxys/proxys.go b/internal/proxys/proxys.go
@@ -18,11 +18,17 @@ func (p *IProxy) GetIPS() int {
 }
 
 func (p *IProxy) GetProxyIP() string {
+	if p == nil {
+		return ""
+	}
+
 	p.lock.Lock()
 	defer p.lock.Unlock()
+
 	if len(p.ips) == 0 {
 		return ""
 	}
+
 	proxyIp := p.ips[0]
 	p.ips = append(p.ips[1:], proxyIp)
 	return proxyIp

diff --git a/internal/tokens/tokens.go b/internal/tokens/tokens.go
@@ -29,6 +29,10 @@ func NewAccessToken(tokens []*Secret) AccessToken {
 }
 
 func (a *AccessToken) GetSecret() *Secret {
+	if a == nil {
+		return &Secret{}
+	}
+
 	a.lock.Lock()
 	defer a.lock.Unlock()