darroyo
f775c2f67e
CI / test (push) Failing after 12m57s
- manager.go: add _ = to syscall.Kill and proc.Signal calls (6 fixes) - chat.go: add _ = to c.Writer.WriteString calls (9 fixes) - webhook.go: remove unused webhookSecret field
107 lines
2.5 KiB
Go
107 lines
2.5 KiB
Go
package quota
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/hmac"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log/slog"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"gorm.io/gorm"
|
|
|
|
"github.com/llamalink/llamalink/internal/db"
|
|
)
|
|
|
|
type WebhookPayload struct {
|
|
Event string `json:"event"`
|
|
Timestamp time.Time `json:"timestamp"`
|
|
KeyID uuid.UUID `json:"key_id"`
|
|
KeyName string `json:"key_name"`
|
|
Message string `json:"message"`
|
|
Usage *struct {
|
|
Used int `json:"tokens_used"`
|
|
Limit int `json:"tokens_limit"`
|
|
} `json:"usage,omitempty"`
|
|
}
|
|
|
|
type WebhookService struct {
|
|
db *gorm.DB
|
|
client *http.Client
|
|
}
|
|
|
|
func NewWebhookService(db *gorm.DB) *WebhookService {
|
|
return &WebhookService{
|
|
db: db,
|
|
client: &http.Client{Timeout: 10 * time.Second},
|
|
}
|
|
}
|
|
|
|
func (s *WebhookService) Dispatch(event string, key *db.ApiKey, message string, usage *struct{ Used, Limit int }) {
|
|
var webhooks []db.Webhook
|
|
s.db.Where("api_key_id = ? AND event = ? AND is_active = ?", key.ID, event, true).Find(&webhooks)
|
|
|
|
for _, wh := range webhooks {
|
|
go s.send(wh, event, key, message, usage)
|
|
}
|
|
}
|
|
|
|
func (s *WebhookService) send(wh db.Webhook, event string, key *db.ApiKey, message string, usage *struct{ Used, Limit int }) {
|
|
payload := WebhookPayload{
|
|
Event: event,
|
|
Timestamp: time.Now().UTC(),
|
|
KeyID: key.ID,
|
|
KeyName: key.Name,
|
|
Message: message,
|
|
}
|
|
|
|
if usage != nil {
|
|
payload.Usage = &struct {
|
|
Used int `json:"tokens_used"`
|
|
Limit int `json:"tokens_limit"`
|
|
}{Used: usage.Used, Limit: usage.Limit}
|
|
}
|
|
|
|
body, err := json.Marshal(payload)
|
|
if err != nil {
|
|
slog.Error("failed to marshal webhook payload", "error", err)
|
|
return
|
|
}
|
|
|
|
req, err := http.NewRequest("POST", wh.URL, bytes.NewReader(body))
|
|
if err != nil {
|
|
slog.Error("failed to create webhook request", "error", err)
|
|
return
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("X-LlamaLink-Event", event)
|
|
req.Header.Set("X-LlamaLink-Timestamp", fmt.Sprintf("%d", time.Now().Unix()))
|
|
|
|
if wh.Secret != nil && *wh.Secret != "" {
|
|
sig := s.sign(body, *wh.Secret)
|
|
req.Header.Set("X-LlamaLink-Signature", sig)
|
|
}
|
|
|
|
resp, err := s.client.Do(req)
|
|
if err != nil {
|
|
slog.Error("webhook delivery failed", "url", wh.URL, "error", err)
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode >= 400 {
|
|
slog.Warn("webhook returned error", "url", wh.URL, "status", resp.StatusCode)
|
|
}
|
|
}
|
|
|
|
func (s *WebhookService) sign(body []byte, secret string) string {
|
|
h := hmac.New(sha256.New, []byte(secret))
|
|
h.Write(body)
|
|
return "sha256=" + hex.EncodeToString(h.Sum(nil))
|
|
}
|