Initial zhanlu OpenAI proxy

internal/auth/sso.go

@@ -0,0 +1,90 @@
+package auth
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"strings"
+	"time"
+)
+
+type ExchangeResponse struct {
+	ErrorCode string         `json:"errorCode"`
+	ErrorMsg  string         `json:"errorMsg"`
+	Message   string         `json:"message"`
+	Body      map[string]any `json:"body"`
+}
+
+func ExchangeCode(client *http.Client, endpoint string, code string, decryptKey string) (Credentials, error) {
+	if strings.TrimSpace(endpoint) == "" {
+		return Credentials{}, errors.New("exchange endpoint is required")
+	}
+	if strings.TrimSpace(code) == "" {
+		return Credentials{}, errors.New("code is required")
+	}
+	if strings.TrimSpace(decryptKey) == "" {
+		return Credentials{}, errors.New("decrypt key is required")
+	}
+	if client == nil {
+		client = &http.Client{Timeout: 60 * time.Second}
+	}
+
+	body, _ := json.Marshal(map[string]string{"code": code})
+	req, err := http.NewRequest(http.MethodPost, endpoint, bytes.NewReader(body))
+	if err != nil {
+		return Credentials{}, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := client.Do(req)
+	if err != nil {
+		return Credentials{}, err
+	}
+	defer resp.Body.Close()
+
+	var exchange ExchangeResponse
+	if err := json.NewDecoder(resp.Body).Decode(&exchange); err != nil {
+		return Credentials{}, err
+	}
+	if exchange.ErrorCode != "Success" {
+		msg := firstNonEmpty(exchange.ErrorMsg, exchange.Message, exchange.ErrorCode)
+		return Credentials{}, fmt.Errorf("exchange failed: %s", msg)
+	}
+
+	ak, err := decryptBodyField(exchange.Body, "ak", decryptKey)
+	if err != nil {
+		return Credentials{}, err
+	}
+	sk, err := decryptBodyField(exchange.Body, "sk", decryptKey)
+	if err != nil {
+		return Credentials{}, err
+	}
+	token, err := decryptBodyField(exchange.Body, "token", decryptKey)
+	if err != nil {
+		return Credentials{}, err
+	}
+	return Credentials{AccessKey: ak, SecretKey: sk, Token: token, SavedAt: time.Now()}, nil
+}
+
+func decryptBodyField(body map[string]any, key string, decryptKey string) (string, error) {
+	v, ok := body[key]
+	if !ok {
+		return "", fmt.Errorf("response body missing %s", key)
+	}
+	s, ok := v.(string)
+	if !ok || strings.TrimSpace(s) == "" {
+		return "", fmt.Errorf("response body %s is not a string", key)
+	}
+	return DecryptCredential(strings.TrimSpace(s), decryptKey)
+}
+
+func firstNonEmpty(values ...string) string {
+	for _, v := range values {
+		if strings.TrimSpace(v) != "" {
+			return v
+		}
+	}
+	return "unknown error"
+}