bedrock_signer.go 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. package service
  2. import (
  3. "context"
  4. "crypto/sha256"
  5. "encoding/hex"
  6. "fmt"
  7. "net/http"
  8. "time"
  9. "github.com/aws/aws-sdk-go-v2/aws"
  10. v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
  11. )
  12. // BedrockSigner 使用 AWS SigV4 对 Bedrock 请求签名
  13. type BedrockSigner struct {
  14. credentials aws.Credentials
  15. region string
  16. signer *v4.Signer
  17. }
  18. // NewBedrockSigner 创建 BedrockSigner
  19. func NewBedrockSigner(accessKeyID, secretAccessKey, sessionToken, region string) *BedrockSigner {
  20. return &BedrockSigner{
  21. credentials: aws.Credentials{
  22. AccessKeyID: accessKeyID,
  23. SecretAccessKey: secretAccessKey,
  24. SessionToken: sessionToken,
  25. },
  26. region: region,
  27. signer: v4.NewSigner(),
  28. }
  29. }
  30. // NewBedrockSignerFromAccount 从 Account 凭证创建 BedrockSigner
  31. func NewBedrockSignerFromAccount(account *Account) (*BedrockSigner, error) {
  32. accessKeyID := account.GetCredential("aws_access_key_id")
  33. if accessKeyID == "" {
  34. return nil, fmt.Errorf("aws_access_key_id not found in credentials")
  35. }
  36. secretAccessKey := account.GetCredential("aws_secret_access_key")
  37. if secretAccessKey == "" {
  38. return nil, fmt.Errorf("aws_secret_access_key not found in credentials")
  39. }
  40. region := account.GetCredential("aws_region")
  41. if region == "" {
  42. region = defaultBedrockRegion
  43. }
  44. sessionToken := account.GetCredential("aws_session_token") // 可选
  45. return NewBedrockSigner(accessKeyID, secretAccessKey, sessionToken, region), nil
  46. }
  47. // SignRequest 对 HTTP 请求进行 SigV4 签名
  48. // 重要约束:调用此方法前,req 应只包含 AWS 相关的 header(如 Content-Type、Accept)。
  49. // 非 AWS header(如 anthropic-beta)会参与签名计算,如果 Bedrock 服务端不识别这些 header,
  50. // 签名验证可能失败。litellm 通过 _filter_headers_for_aws_signature 实现头过滤,
  51. // 当前实现中 buildUpstreamRequestBedrock 仅设置了 Content-Type 和 Accept,因此是安全的。
  52. func (s *BedrockSigner) SignRequest(ctx context.Context, req *http.Request, body []byte) error {
  53. payloadHash := sha256Hash(body)
  54. return s.signer.SignHTTP(ctx, s.credentials, req, payloadHash, "bedrock", s.region, time.Now())
  55. }
  56. func sha256Hash(data []byte) string {
  57. h := sha256.Sum256(data)
  58. return hex.EncodeToString(h[:])
  59. }