This commit is contained in:
Ylarod
2026-03-13 19:15:27 +08:00
parent 11f7b83522
commit e90ec847b6
5 changed files with 37 additions and 56 deletions

View File

@@ -27,10 +27,10 @@ const (
// Account type constants // Account type constants
const ( const (
AccountTypeOAuth = "oauth" // OAuth类型账号full scope: profile + inference AccountTypeOAuth = "oauth" // OAuth类型账号full scope: profile + inference
AccountTypeSetupToken = "setup-token" // Setup Token类型账号inference only scope AccountTypeSetupToken = "setup-token" // Setup Token类型账号inference only scope
AccountTypeAPIKey = "apikey" // API Key类型账号 AccountTypeAPIKey = "apikey" // API Key类型账号
AccountTypeUpstream = "upstream" // 上游透传类型账号(通过 Base URL + API Key 连接上游) AccountTypeUpstream = "upstream" // 上游透传类型账号(通过 Base URL + API Key 连接上游)
AccountTypeBedrock = "bedrock" // AWS Bedrock 类型账号(通过 SigV4 签名连接 Bedrock AccountTypeBedrock = "bedrock" // AWS Bedrock 类型账号(通过 SigV4 签名连接 Bedrock
AccountTypeBedrockAPIKey = "bedrock-apikey" // AWS Bedrock API Key 类型账号(通过 Bearer Token 连接 Bedrock AccountTypeBedrockAPIKey = "bedrock-apikey" // AWS Bedrock API Key 类型账号(通过 Bearer Token 连接 Bedrock
) )

View File

@@ -282,8 +282,8 @@ func (d *bedrockEventStreamDecoder) Decode() ([]byte, error) {
// 验证 message CRC覆盖 prelude + headers + payload // 验证 message CRC覆盖 prelude + headers + payload
messageCRC := bedrockReadUint32(data[len(data)-4:]) messageCRC := bedrockReadUint32(data[len(data)-4:])
h := crc32.New(crc32IEEETable) h := crc32.New(crc32IEEETable)
h.Write(prelude) _, _ = h.Write(prelude)
h.Write(data[:len(data)-4]) _, _ = h.Write(data[:len(data)-4])
if h.Sum32() != messageCRC { if h.Sum32() != messageCRC {
return nil, fmt.Errorf("eventstream message CRC mismatch") return nil, fmt.Errorf("eventstream message CRC mismatch")
} }

View File

@@ -76,15 +76,15 @@ func TestExtractEventStreamHeaderValue(t *testing.T) {
buildStringHeader := func(name, value string) []byte { buildStringHeader := func(name, value string) []byte {
var buf bytes.Buffer var buf bytes.Buffer
// name length (1 byte) // name length (1 byte)
buf.WriteByte(byte(len(name))) _ = buf.WriteByte(byte(len(name)))
// name // name
buf.WriteString(name) _, _ = buf.WriteString(name)
// value type (7 = string) // value type (7 = string)
buf.WriteByte(7) _ = buf.WriteByte(7)
// value length (2 bytes, big-endian) // value length (2 bytes, big-endian)
_ = binary.Write(&buf, binary.BigEndian, uint16(len(value))) _ = binary.Write(&buf, binary.BigEndian, uint16(len(value)))
// value // value
buf.WriteString(value) _, _ = buf.WriteString(value)
return buf.Bytes() return buf.Bytes()
} }
@@ -100,9 +100,9 @@ func TestExtractEventStreamHeaderValue(t *testing.T) {
t.Run("multiple headers", func(t *testing.T) { t.Run("multiple headers", func(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
buf.Write(buildStringHeader(":content-type", "application/json")) _, _ = buf.Write(buildStringHeader(":content-type", "application/json"))
buf.Write(buildStringHeader(":event-type", "chunk")) _, _ = buf.Write(buildStringHeader(":event-type", "chunk"))
buf.Write(buildStringHeader(":message-type", "event")) _, _ = buf.Write(buildStringHeader(":message-type", "event"))
headers := buf.Bytes() headers := buf.Bytes()
assert.Equal(t, "chunk", extractEventStreamHeaderValue(headers, ":event-type")) assert.Equal(t, "chunk", extractEventStreamHeaderValue(headers, ":event-type"))
@@ -123,17 +123,17 @@ func TestBedrockEventStreamDecoder(t *testing.T) {
// Build headers // Build headers
var headersBuf bytes.Buffer var headersBuf bytes.Buffer
// :event-type header // :event-type header
headersBuf.WriteByte(byte(len(":event-type"))) _ = headersBuf.WriteByte(byte(len(":event-type")))
headersBuf.WriteString(":event-type") _, _ = headersBuf.WriteString(":event-type")
headersBuf.WriteByte(7) // string type _ = headersBuf.WriteByte(7) // string type
_ = binary.Write(&headersBuf, binary.BigEndian, uint16(len(eventType))) _ = binary.Write(&headersBuf, binary.BigEndian, uint16(len(eventType)))
headersBuf.WriteString(eventType) _, _ = headersBuf.WriteString(eventType)
// :message-type header // :message-type header
headersBuf.WriteByte(byte(len(":message-type"))) _ = headersBuf.WriteByte(byte(len(":message-type")))
headersBuf.WriteString(":message-type") _, _ = headersBuf.WriteString(":message-type")
headersBuf.WriteByte(7) _ = headersBuf.WriteByte(7)
_ = binary.Write(&headersBuf, binary.BigEndian, uint16(len("event"))) _ = binary.Write(&headersBuf, binary.BigEndian, uint16(len("event")))
headersBuf.WriteString("event") _, _ = headersBuf.WriteString("event")
headers := headersBuf.Bytes() headers := headersBuf.Bytes()
headersLen := uint32(len(headers)) headersLen := uint32(len(headers))
@@ -149,10 +149,10 @@ func TestBedrockEventStreamDecoder(t *testing.T) {
// Build frame: prelude + prelude_crc + headers + payload // Build frame: prelude + prelude_crc + headers + payload
var frame bytes.Buffer var frame bytes.Buffer
frame.Write(preludeBytes) _, _ = frame.Write(preludeBytes)
_ = binary.Write(&frame, binary.BigEndian, preludeCRC) _ = binary.Write(&frame, binary.BigEndian, preludeCRC)
frame.Write(headers) _, _ = frame.Write(headers)
frame.Write(payload) _, _ = frame.Write(payload)
// Message CRC covers everything before itself // Message CRC covers everything before itself
messageCRC := crc32.Checksum(frame.Bytes(), crc32IeeeTab) messageCRC := crc32.Checksum(frame.Bytes(), crc32IeeeTab)
@@ -173,9 +173,9 @@ func TestBedrockEventStreamDecoder(t *testing.T) {
t.Run("skip non-chunk events", func(t *testing.T) { t.Run("skip non-chunk events", func(t *testing.T) {
// Write initial-response followed by chunk // Write initial-response followed by chunk
var buf bytes.Buffer var buf bytes.Buffer
buf.Write(buildFrame("initial-response", []byte(`{}`))) _, _ = buf.Write(buildFrame("initial-response", []byte(`{}`)))
chunkPayload := []byte(`{"bytes":"aGVsbG8="}`) chunkPayload := []byte(`{"bytes":"aGVsbG8="}`)
buf.Write(buildFrame("chunk", chunkPayload)) _, _ = buf.Write(buildFrame("chunk", chunkPayload))
decoder := newBedrockEventStreamDecoder(&buf) decoder := newBedrockEventStreamDecoder(&buf)
result, err := decoder.Decode() result, err := decoder.Decode()
@@ -214,11 +214,11 @@ func TestBedrockEventStreamDecoder(t *testing.T) {
payload := []byte(`{"bytes":"dGVzdA=="}`) payload := []byte(`{"bytes":"dGVzdA=="}`)
var headersBuf bytes.Buffer var headersBuf bytes.Buffer
headersBuf.WriteByte(byte(len(":event-type"))) _ = headersBuf.WriteByte(byte(len(":event-type")))
headersBuf.WriteString(":event-type") _, _ = headersBuf.WriteString(":event-type")
headersBuf.WriteByte(7) _ = headersBuf.WriteByte(7)
_ = binary.Write(&headersBuf, binary.BigEndian, uint16(len("chunk"))) _ = binary.Write(&headersBuf, binary.BigEndian, uint16(len("chunk")))
headersBuf.WriteString("chunk") _, _ = headersBuf.WriteString("chunk")
headers := headersBuf.Bytes() headers := headersBuf.Bytes()
headersLen := uint32(len(headers)) headersLen := uint32(len(headers))
@@ -230,10 +230,10 @@ func TestBedrockEventStreamDecoder(t *testing.T) {
preludeBytes := preludeBuf.Bytes() preludeBytes := preludeBuf.Bytes()
var frame bytes.Buffer var frame bytes.Buffer
frame.Write(preludeBytes) _, _ = frame.Write(preludeBytes)
_ = binary.Write(&frame, binary.BigEndian, crc32.Checksum(preludeBytes, castagnoliTab)) _ = binary.Write(&frame, binary.BigEndian, crc32.Checksum(preludeBytes, castagnoliTab))
frame.Write(headers) _, _ = frame.Write(headers)
frame.Write(payload) _, _ = frame.Write(payload)
_ = binary.Write(&frame, binary.BigEndian, crc32.Checksum(frame.Bytes(), castagnoliTab)) _ = binary.Write(&frame, binary.BigEndian, crc32.Checksum(frame.Bytes(), castagnoliTab))
decoder := newBedrockEventStreamDecoder(bytes.NewReader(frame.Bytes())) decoder := newBedrockEventStreamDecoder(bytes.NewReader(frame.Bytes()))

View File

@@ -29,10 +29,10 @@ const (
// Account type constants // Account type constants
const ( const (
AccountTypeOAuth = domain.AccountTypeOAuth // OAuth类型账号full scope: profile + inference AccountTypeOAuth = domain.AccountTypeOAuth // OAuth类型账号full scope: profile + inference
AccountTypeSetupToken = domain.AccountTypeSetupToken // Setup Token类型账号inference only scope AccountTypeSetupToken = domain.AccountTypeSetupToken // Setup Token类型账号inference only scope
AccountTypeAPIKey = domain.AccountTypeAPIKey // API Key类型账号 AccountTypeAPIKey = domain.AccountTypeAPIKey // API Key类型账号
AccountTypeUpstream = domain.AccountTypeUpstream // 上游透传类型账号(通过 Base URL + API Key 连接上游) AccountTypeUpstream = domain.AccountTypeUpstream // 上游透传类型账号(通过 Base URL + API Key 连接上游)
AccountTypeBedrock = domain.AccountTypeBedrock // AWS Bedrock 类型账号(通过 SigV4 签名连接 Bedrock AccountTypeBedrock = domain.AccountTypeBedrock // AWS Bedrock 类型账号(通过 SigV4 签名连接 Bedrock
AccountTypeBedrockAPIKey = domain.AccountTypeBedrockAPIKey // AWS Bedrock API Key 类型账号(通过 Bearer Token 连接 Bedrock AccountTypeBedrockAPIKey = domain.AccountTypeBedrockAPIKey // AWS Bedrock API Key 类型账号(通过 Bearer Token 连接 Bedrock
) )

View File

@@ -5853,25 +5853,6 @@ func containsBetaToken(header, token string) bool {
return false return false
} }
// filterBetaTokensFromHeader removes tokens present in filterSet from a comma-separated header value.
// Returns the filtered header string, or "" if all tokens were removed.
func filterBetaTokensFromHeader(header string, filterSet map[string]struct{}) string {
if header == "" || len(filterSet) == 0 {
return header
}
var kept []string
for _, p := range strings.Split(header, ",") {
t := strings.TrimSpace(p)
if t == "" {
continue
}
if _, filtered := filterSet[t]; !filtered {
kept = append(kept, t)
}
}
return strings.Join(kept, ", ")
}
func filterBetaTokens(tokens []string, filterSet map[string]struct{}) []string { func filterBetaTokens(tokens []string, filterSet map[string]struct{}) []string {
if len(tokens) == 0 || len(filterSet) == 0 { if len(tokens) == 0 || len(filterSet) == 0 {
return tokens return tokens