78 lines
2.1 KiB
Go
78 lines
2.1 KiB
Go
package providerapi
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"math"
|
|
)
|
|
|
|
var (
|
|
ErrResponseTooLarge = errors.New("provider response exceeds byte limit")
|
|
ErrTemplateInputTooLarge = errors.New("provider template input exceeds byte limit")
|
|
ErrTemplateOutputTooLarge = errors.New("provider template output exceeds byte limit")
|
|
ErrTooManyCandidates = errors.New("provider template output exceeds candidate limit")
|
|
ErrInvalidProxyOutput = errors.New("provider template output contains no valid proxy")
|
|
ErrRecursiveTemplate = errors.New("provider template contains recursive calls")
|
|
ErrTemplateTooComplex = errors.New("provider template exceeds static complexity limit")
|
|
ErrCredentialStoreRequired = errors.New("provider credentials require a credential store")
|
|
ErrInvalidHTTPResponse = errors.New("provider HTTP client returned an invalid response")
|
|
)
|
|
|
|
type limitError struct {
|
|
kind error
|
|
size int64
|
|
limit int64
|
|
}
|
|
|
|
func (e *limitError) Error() string {
|
|
return fmt.Sprintf("%v: size %d, limit %d", e.kind, e.size, e.limit)
|
|
}
|
|
|
|
func (e *limitError) Unwrap() error { return e.kind }
|
|
|
|
func enforceByteLimit(kind error, size, limit int64) error {
|
|
if size <= limit {
|
|
return nil
|
|
}
|
|
return &limitError{kind: kind, size: size, limit: limit}
|
|
}
|
|
|
|
func readAllLimited(reader io.Reader, limit int64, kind error) ([]byte, error) {
|
|
readLimit := limit
|
|
if readLimit < math.MaxInt64 {
|
|
readLimit++
|
|
}
|
|
limited := &io.LimitedReader{R: reader, N: readLimit}
|
|
data, err := io.ReadAll(limited)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if err := enforceByteLimit(kind, int64(len(data)), limit); err != nil {
|
|
return nil, err
|
|
}
|
|
return data, nil
|
|
}
|
|
|
|
type limitedBuffer struct {
|
|
ctx context.Context
|
|
kind error
|
|
limit int64
|
|
buffer bytes.Buffer
|
|
}
|
|
|
|
func (w *limitedBuffer) Write(data []byte) (int, error) {
|
|
if err := w.ctx.Err(); err != nil {
|
|
return 0, err
|
|
}
|
|
current := int64(w.buffer.Len())
|
|
if err := enforceByteLimit(w.kind, current+int64(len(data)), w.limit); err != nil {
|
|
return 0, err
|
|
}
|
|
return w.buffer.Write(data)
|
|
}
|
|
|
|
func (w *limitedBuffer) String() string { return w.buffer.String() }
|