Initial zhanlu OpenAI proxy

internal/auth/aes.go

@@ -0,0 +1,85 @@
+package auth
+
+import (
+	"crypto/aes"
+	"encoding/base64"
+	"errors"
+	"fmt"
+)
+
+func DecryptCredential(ciphertextBase64, key string) (string, error) {
+	ciphertext, err := base64.StdEncoding.DecodeString(ciphertextBase64)
+	if err != nil {
+		return "", err
+	}
+	if len(ciphertext) == 0 || len(ciphertext)%aes.BlockSize != 0 {
+		return "", errors.New("invalid AES ciphertext length")
+	}
+
+	block, err := aes.NewCipher(repeatKey(key, aes.BlockSize))
+	if err != nil {
+		return "", err
+	}
+
+	plain := make([]byte, len(ciphertext))
+	for start := 0; start < len(ciphertext); start += aes.BlockSize {
+		block.Decrypt(plain[start:start+aes.BlockSize], ciphertext[start:start+aes.BlockSize])
+	}
+
+	plain, err = unpadPKCS7(plain, aes.BlockSize)
+	if err != nil {
+		return "", err
+	}
+	return string(plain), nil
+}
+
+func EncryptCredential(plaintext, key string) (string, error) {
+	block, err := aes.NewCipher(repeatKey(key, aes.BlockSize))
+	if err != nil {
+		return "", err
+	}
+	plain := padPKCS7([]byte(plaintext), aes.BlockSize)
+	out := make([]byte, len(plain))
+	for start := 0; start < len(plain); start += aes.BlockSize {
+		block.Encrypt(out[start:start+aes.BlockSize], plain[start:start+aes.BlockSize])
+	}
+	return base64.StdEncoding.EncodeToString(out), nil
+}
+
+func repeatKey(key string, size int) []byte {
+	src := []byte(key)
+	out := make([]byte, size)
+	if len(src) == 0 {
+		return out
+	}
+	for i := range out {
+		out[i] = src[i%len(src)]
+	}
+	return out
+}
+
+func unpadPKCS7(in []byte, blockSize int) ([]byte, error) {
+	if len(in) == 0 || len(in)%blockSize != 0 {
+		return nil, errors.New("invalid padded data length")
+	}
+	pad := int(in[len(in)-1])
+	if pad == 0 || pad > blockSize || pad > len(in) {
+		return nil, fmt.Errorf("invalid padding size %d", pad)
+	}
+	for _, b := range in[len(in)-pad:] {
+		if int(b) != pad {
+			return nil, errors.New("invalid padding bytes")
+		}
+	}
+	return in[:len(in)-pad], nil
+}
+
+func padPKCS7(in []byte, blockSize int) []byte {
+	pad := blockSize - len(in)%blockSize
+	out := make([]byte, len(in)+pad)
+	copy(out, in)
+	for i := len(in); i < len(out); i++ {
+		out[i] = byte(pad)
+	}
+	return out
+}

internal/auth/credentials.go

@@ -0,0 +1,58 @@
+package auth
+
+import (
+	"encoding/json"
+	"errors"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+)
+
+type Credentials struct {
+	AccessKey string    `json:"access_key"`
+	SecretKey string    `json:"secret_key"`
+	Token     string    `json:"token"`
+	BaseURL   string    `json:"base_url,omitempty"`
+	SavedAt   time.Time `json:"saved_at"`
+}
+
+func (c Credentials) Validate() error {
+	if strings.TrimSpace(c.AccessKey) == "" {
+		return errors.New("access_key is required")
+	}
+	if strings.TrimSpace(c.SecretKey) == "" {
+		return errors.New("secret_key is required")
+	}
+	if strings.TrimSpace(c.Token) == "" {
+		return errors.New("token is required")
+	}
+	return nil
+}
+
+func LoadCredentials(path string) (Credentials, error) {
+	b, err := os.ReadFile(path)
+	if err != nil {
+		return Credentials{}, err
+	}
+	var c Credentials
+	if err := json.Unmarshal(b, &c); err != nil {
+		return Credentials{}, err
+	}
+	return c, c.Validate()
+}
+
+func SaveCredentials(path string, c Credentials) error {
+	if err := c.Validate(); err != nil {
+		return err
+	}
+	c.SavedAt = time.Now()
+	b, err := json.MarshalIndent(c, "", "  ")
+	if err != nil {
+		return err
+	}
+	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
+		return err
+	}
+	return os.WriteFile(path, b, 0o600)
+}

internal/auth/rsa.go

@@ -0,0 +1,37 @@
+package auth
+
+import (
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/base64"
+	"encoding/pem"
+	"errors"
+)
+
+func ParsePublicKey(pemText string) (*rsa.PublicKey, error) {
+	block, _ := pem.Decode([]byte(pemText))
+	if block == nil {
+		return nil, errors.New("invalid public key PEM")
+	}
+	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
+	if err != nil {
+		return nil, err
+	}
+	rsaPub, ok := pub.(*rsa.PublicKey)
+	if !ok {
+		return nil, errors.New("public key is not RSA")
+	}
+	return rsaPub, nil
+}
+
+func EncryptAuthorization(pub *rsa.PublicKey, plaintext string) (string, error) {
+	if pub == nil {
+		return "", errors.New("public key is required")
+	}
+	out, err := rsa.EncryptPKCS1v15(rand.Reader, pub, []byte(plaintext))
+	if err != nil {
+		return "", err
+	}
+	return base64.StdEncoding.EncodeToString(out), nil
+}

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"
+}