70 lines
2.0 KiB
Go
70 lines
2.0 KiB
Go
package memorytool
|
|
|
|
import (
|
|
"crypto/hmac"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"errors"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
var ErrInvalidToken = errors.New("invalid memory tool token")
|
|
|
|
type Claims struct {
|
|
ConversationID string `json:"conversation_id"`
|
|
TaskRef string `json:"task_ref"`
|
|
ExpiresAt int64 `json:"expires_at"`
|
|
}
|
|
|
|
type TokenSigner struct {
|
|
secret []byte
|
|
now func() time.Time
|
|
}
|
|
|
|
func NewTokenSigner(secret []byte) (*TokenSigner, error) {
|
|
if len(secret) < 32 {
|
|
return nil, errors.New("memory tool signing secret must contain at least 32 bytes")
|
|
}
|
|
return &TokenSigner{secret: append([]byte(nil), secret...), now: time.Now}, nil
|
|
}
|
|
|
|
func (s *TokenSigner) Issue(conversationID, taskRef string, expiresAt time.Time) (string, error) {
|
|
if strings.TrimSpace(conversationID) == "" || strings.TrimSpace(taskRef) == "" || !expiresAt.After(s.now()) {
|
|
return "", ErrInvalidToken
|
|
}
|
|
payload, err := json.Marshal(Claims{ConversationID: conversationID, TaskRef: taskRef, ExpiresAt: expiresAt.UTC().Unix()})
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
encoded := base64.RawURLEncoding.EncodeToString(payload)
|
|
return "v1." + encoded + "." + s.signature(encoded), nil
|
|
}
|
|
|
|
func (s *TokenSigner) Verify(token string) (Claims, error) {
|
|
parts := strings.Split(token, ".")
|
|
if len(parts) != 3 || parts[0] != "v1" {
|
|
return Claims{}, ErrInvalidToken
|
|
}
|
|
expected := s.signature(parts[1])
|
|
if !hmac.Equal([]byte(expected), []byte(parts[2])) {
|
|
return Claims{}, ErrInvalidToken
|
|
}
|
|
payload, err := base64.RawURLEncoding.DecodeString(parts[1])
|
|
if err != nil {
|
|
return Claims{}, ErrInvalidToken
|
|
}
|
|
var claims Claims
|
|
if err := json.Unmarshal(payload, &claims); err != nil || claims.ConversationID == "" || claims.TaskRef == "" || claims.ExpiresAt <= s.now().UTC().Unix() {
|
|
return Claims{}, ErrInvalidToken
|
|
}
|
|
return claims, nil
|
|
}
|
|
|
|
func (s *TokenSigner) signature(payload string) string {
|
|
mac := hmac.New(sha256.New, s.secret)
|
|
_, _ = mac.Write([]byte("memory-tool-token-v1." + payload))
|
|
return base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
|
|
}
|