opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(security,worker,agent,pcerrors): add security hardening + job worker + agent state

Browse source 
security:
- jsonextract.go — hardened JSON extraction against prompt injection;
  multiple fallback strategies (regex, manual brace-scan, strip-think)
- redact.go — auto-strip API keys, auth headers, PII from log strings
- urlguard.go — SSRF blocklist (RFC1918 + link-local + metadata IPs)
- vault.go — ChaCha20-Poly1305 encrypted secret storage with keyring

worker:
- worker.go — background job processor polling the jobs table;
  claim-via-UPDATE, exponential backoff, max_attempts enforcement

agent:
- kv.go — agent KV store accessor wrapping delegate.UpsertKV/GetKV
- state_store.go — conversation + run + run_state persistence layer
- offloading_tool_runtime.go — tool result offload coordinator;
  ShouldOffload threshold check + OffloadToolResult flow
- tool_result_search.go — fetch offloaded tool results by run/step

pcerrors:
- pcerrors.go — errbuilder-based domain error types for picoclaw
- cli.go — human-friendly CLI error rendering
This commit is contained in:
ZanzyTHEbar 2026-02-18 23:41:33 +00:00
parent 4b1567f08b
commit 75c9e84309
15 changed files with 2374 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
pkg/agent/kv.go Normal file
View file

@ -0,0 +1,13 @@
package agent
import "context"
// KVDelegate is a minimal KV-like interface for offloading large artifacts
// (tool results, scratchpads) out of the LLM context window.
//
// Keys are logical paths (not filesystem paths).
type KVDelegate interface {
Put(ctx context.Context, key string, value []byte) error
Get(ctx context.Context, key string) ([]byte, error)
Scan(ctx context.Context, prefix string) ([]string, error)
}

277
pkg/agent/offloading_tool_runtime.go Normal file
View file

@ -0,0 +1,277 @@
package agent
import (
"context"
"encoding/json"
"fmt"
"strconv"
"strings"
"charm.land/fantasy"
"github.com/sipeed/picoclaw/pkg/ids"
"github.com/sipeed/picoclaw/pkg/memory/sqlc"
"github.com/sipeed/picoclaw/pkg/pcerrors"
)
const defaultToolMaxConcurrency = 4
type ctxStepIndexKey struct{}
func WithStepIndex(ctx context.Context, stepIndex int) context.Context {
return context.WithValue(ctx, ctxStepIndexKey{}, stepIndex)
}
func StepIndexFromCtx(ctx context.Context) int {
v := ctx.Value(ctxStepIndexKey{})
if v == nil {
return 0
}
if i, ok := v.(int); ok {
return i
}
return 0
}
// OffloadingToolRuntime wraps a base ToolRuntime and applies tool result
// offloading policy:
// - Always offload full results to KV delegate.
// - If result is below threshold, keep it inline as-is.
// - If above threshold, truncate inline output and include an index/instructions.
// - Store chunked payload for targeted retrieval.
type OffloadingToolRuntime struct {
Base fantasy.ToolRuntime
KV KVDelegate
Queries *sqlc.Queries
ConversationID ids.UUID
RunID ids.UUID
ThresholdChars int
ChunkChars int
}
func (r OffloadingToolRuntime) Execute(ctx context.Context, tools []fantasy.AgentTool, toolCalls []fantasy.ToolCallContent, _ func(result fantasy.ToolResultContent) error) ([]fantasy.ToolResultContent, error) {
if len(toolCalls) == 0 {
return nil, nil
}
if r.Base == nil {
r.Base = fantasy.DAGToolRuntime{MaxConcurrency: defaultToolMaxConcurrency}
}
if r.KV == nil {
return nil, pcerrors.New(pcerrors.CodeFailedPrecondition, "KV delegate is nil")
}
if r.Queries == nil {
return nil, pcerrors.New(pcerrors.CodeFailedPrecondition, "db queries is nil")
}
if r.ConversationID.IsZero() || r.RunID.IsZero() {
return nil, pcerrors.New(pcerrors.CodeInvalidArgument, "conversation_id/run_id is required")
}
threshold := r.ThresholdChars
if threshold <= 0 {
threshold = 4_000
}
chunkChars := r.ChunkChars
if chunkChars <= 0 {
chunkChars = 2_000
}
stepIndex := StepIndexFromCtx(ctx)
results, err := r.Base.Execute(ctx, tools, toolCalls, nil)
if err != nil {
return nil, err
}
for i := range results {
tc := toolCalls[i]
res := results[i]
fullKey := toolResultFullKey(r.ConversationID, r.RunID, stepIndex, tc.ToolCallID)
payload, payloadType, payloadText := toolResultPayload(res)
b, marshalErr := json.Marshal(payload)
if marshalErr != nil {
b = []byte(`{"error":"failed to marshal tool result payload"}`)
payloadType = "error"
}
if putErr := r.KV.Put(ctx, fullKey, b); putErr != nil {
return nil, putErr
}
preview := payloadText
chunkCount := int64(0)
if payloadType == "text" && len([]rune(payloadText)) > threshold {
chunks := chunkString(payloadText, chunkChars)
chunkCount = int64(len(chunks))
for ci, chunk := range chunks {
chunkKey := toolResultChunkKey(r.ConversationID, r.RunID, stepIndex, tc.ToolCallID, ci)
if putErr := r.KV.Put(ctx, chunkKey, []byte(chunk)); putErr != nil {
return nil, putErr
}
}
preview = truncateRunes(payloadText, threshold)
preview = strings.TrimSpace(preview) + "\n\n" +
"[TRUNCATED]\n" +
"- run_id: " + r.RunID.String() + "\n" +
"- tool_call_id: " + tc.ToolCallID + "\n" +
"- chunk_count: " + strconv.FormatInt(chunkCount, 10) + "\n" +
"Use tool_result_search to retrieve more (prefer chunk ranges)."
results[i].Result = fantasy.ToolResultOutputContentText{Text: preview}
}
if len([]rune(preview)) > 8_000 {
preview = truncateRunes(preview, 8_000)
}
meta := map[string]any{
"result_type": payloadType,
"step_index": stepIndex,
}
metaJSON, _ := json.Marshal(meta)
var previewPtr *string
if strings.TrimSpace(preview) != "" {
previewPtr = &preview
}
_, dbErr := r.Queries.AddAgentToolResult(ctx, sqlc.AddAgentToolResultParams{
ID: ids.New(),
ConversationID: r.ConversationID,
RunID: r.RunID,
StepIndex: int64(stepIndex),
ToolCallID: tc.ToolCallID,
ToolName: tc.ToolName,
FullKey: fullKey,
Preview: previewPtr,
ChunkCount: chunkCount,
MetadataJson: metaJSON,
})
if dbErr != nil {
return nil, dbErr
}
}
return results, nil
}
func toolResultFullKey(conversationID, runID ids.UUID, stepIndex int, toolCallID string) string {
return "tool_results/" + conversationID.String() + "/" + runID.String() + "/step_" + strconv.Itoa(stepIndex) + "/" + sanitizeKeyPart(toolCallID) + "/full.json"
}
func toolResultChunkKey(conversationID, runID ids.UUID, stepIndex int, toolCallID string, chunkIndex int) string {
return "tool_results/" + conversationID.String() + "/" + runID.String() + "/step_" + strconv.Itoa(stepIndex) + "/" + sanitizeKeyPart(toolCallID) + "/chunks/" + fmt.Sprintf("%06d.txt", chunkIndex)
}
func sanitizeKeyPart(s string) string {
s = strings.TrimSpace(s)
if s == "" {
return "empty"
}
var b strings.Builder
b.Grow(len(s))
for _, r := range s {
switch {
case r >= 'a' && r <= 'z':
b.WriteRune(r)
case r >= 'A' && r <= 'Z':
b.WriteRune(r)
case r >= '0' && r <= '9':
b.WriteRune(r)
case r == '-' || r == '_' || r == '.':
b.WriteRune(r)
default:
b.WriteByte('_')
}
}
out := strings.Trim(b.String(), "._")
if out == "" {
return "empty"
}
return out
}
func toolResultPayload(res fantasy.ToolResultContent) (payload map[string]any, payloadType string, payloadText string) {
payloadType = "unknown"
payloadText = ""
switch v := res.Result.(type) {
case fantasy.ToolResultOutputContentText:
payloadType = "text"
payloadText = v.Text
payload = map[string]any{
"type": "text",
"text": v.Text,
}
case fantasy.ToolResultOutputContentMedia:
payloadType = "media"
payloadText = v.Text
payload = map[string]any{
"type": "media",
"text": v.Text,
"media_type": v.MediaType,
"data": v.Data,
}
case fantasy.ToolResultOutputContentError:
payloadType = "error"
errS := ""
if v.Error != nil {
errS = v.Error.Error()
}
payloadText = errS
payload = map[string]any{
"type": "error",
"error": errS,
}
default:
payload = map[string]any{
"type": "unknown",
"value": fmt.Sprintf("%v", res.Result),
}
payloadText = fmt.Sprintf("%v", res.Result)
}
payload["tool_call_id"] = res.ToolCallID
payload["tool_name"] = res.ToolName
payload["provider_executed"] = res.ProviderExecuted
payload["client_metadata"] = res.ClientMetadata
payload["provider_metadata"] = res.ProviderMetadata
return payload, payloadType, payloadText
}
func chunkString(s string, chunkSize int) []string {
if chunkSize <= 0 {
return []string{s}
}
r := []rune(s)
if len(r) == 0 {
return []string{""}
}
out := make([]string, 0, (len(r)/chunkSize)+1)
for i := 0; i < len(r); i += chunkSize {
end := i + chunkSize
if end > len(r) {
end = len(r)
}
out = append(out, string(r[i:end]))
}
return out
}
func truncateRunes(s string, max int) string {
if max <= 0 {
return ""
}
r := []rune(s)
if len(r) <= max {
return s
}
return string(r[:max]) + "…"
}

195
pkg/agent/state_store.go Normal file
View file

@ -0,0 +1,195 @@
package agent
import (
"context"
"encoding/json"
"strings"
"time"
"charm.land/fantasy"
"github.com/sipeed/picoclaw/pkg/ids"
"github.com/sipeed/picoclaw/pkg/memory/sqlc"
"github.com/sipeed/picoclaw/pkg/pcerrors"
)
// StateStore persists agent run state snapshots and transition logs.
type StateStore struct {
q *sqlc.Queries
}
func NewStateStore(q *sqlc.Queries) *StateStore {
return &StateStore{q: q}
}
func (s *StateStore) CreateRun(ctx context.Context, conversationID ids.UUID) (sqlc.AgentRun, error) {
if s == nil || s.q == nil {
return sqlc.AgentRun{}, pcerrors.New(pcerrors.CodeUnknown, "state store is not configured")
}
if conversationID.IsZero() {
return sqlc.AgentRun{}, pcerrors.New(pcerrors.CodeUnknown, "conversation id is empty")
}
return s.q.CreateAgentRun(ctx, sqlc.CreateAgentRunParams{
ID: ids.New(),
ConversationID: conversationID,
Status: "running",
MetadataJson: json.RawMessage(`{}`),
})
}
func (s *StateStore) UpdateRunStatus(ctx context.Context, runID ids.UUID, status string, meta map[string]any) (sqlc.AgentRun, error) {
if s == nil || s.q == nil {
return sqlc.AgentRun{}, pcerrors.New(pcerrors.CodeUnknown, "state store is not configured")
}
if runID.IsZero() {
return sqlc.AgentRun{}, pcerrors.New(pcerrors.CodeUnknown, "run id is empty")
}
if status == "" {
return sqlc.AgentRun{}, pcerrors.New(pcerrors.CodeUnknown, "status is empty")
}
metaJSON := json.RawMessage(`{}`)
if meta != nil {
if b, err := json.Marshal(meta); err == nil {
metaJSON = b
}
}
return s.q.UpdateAgentRunStatus(ctx, sqlc.UpdateAgentRunStatusParams{
Status: status,
MetadataJson: metaJSON,
ID: runID,
})
}
func (s *StateStore) AddRunState(ctx context.Context, runID ids.UUID, stepIndex int, state fantasy.ReActState, snapshot any) (sqlc.AgentRunState, error) {
if s == nil || s.q == nil {
return sqlc.AgentRunState{}, pcerrors.New(pcerrors.CodeUnknown, "state store is not configured")
}
if runID.IsZero() {
return sqlc.AgentRunState{}, pcerrors.New(pcerrors.CodeUnknown, "run id is empty")
}
if stepIndex < 0 {
return sqlc.AgentRunState{}, pcerrors.New(pcerrors.CodeUnknown, "step index is negative")
}
snapJSON := json.RawMessage(`{}`)
if snapshot != nil {
if b, err := json.Marshal(snapshot); err == nil {
snapJSON = b
}
}
return s.q.AddAgentRunState(ctx, sqlc.AddAgentRunStateParams{
ID: ids.New(),
RunID: runID,
StepIndex: int64(stepIndex),
State: string(state),
SnapshotJson: snapJSON,
})
}
func (s *StateStore) AddTransition(ctx context.Context, runID ids.UUID, t fantasy.ReActTransition) (sqlc.AgentStateTransition, error) {
if s == nil || s.q == nil {
return sqlc.AgentStateTransition{}, pcerrors.New(pcerrors.CodeUnknown, "state store is not configured")
}
if runID.IsZero() {
return sqlc.AgentStateTransition{}, pcerrors.New(pcerrors.CodeUnknown, "run id is empty")
}
metaJSON := json.RawMessage(`{}`)
if t.Meta != nil {
if b, err := json.Marshal(t.Meta); err == nil {
metaJSON = b
}
}
var errPtr *string
if t.Error != "" {
errPtr = &t.Error
}
at := t.At
if at.IsZero() {
at = time.Now().UTC()
}
return s.q.AddAgentStateTransition(ctx, sqlc.AddAgentStateTransitionParams{
ID: ids.New(),
RunID: runID,
StepIndex: int64(t.StepIndex),
FromState: string(t.From),
ToState: string(t.To),
Trigger: string(t.Trigger),
At: at,
MetaJson: metaJSON,
Error: errPtr,
})
}
// CheckpointStore persists named checkpoints for later restore.
type CheckpointStore struct {
q *sqlc.Queries
}
func NewCheckpointStore(q *sqlc.Queries) *CheckpointStore {
return &CheckpointStore{q: q}
}
func (s *CheckpointStore) CreateCheckpoint(ctx context.Context, conversationID ids.UUID, name string, runStateID ids.UUID, meta map[string]any) (sqlc.AgentCheckpoint, error) {
if s == nil || s.q == nil {
return sqlc.AgentCheckpoint{}, pcerrors.New(pcerrors.CodeUnknown, "checkpoint store is not configured")
}
if conversationID.IsZero() {
return sqlc.AgentCheckpoint{}, pcerrors.New(pcerrors.CodeUnknown, "conversation id is empty")
}
if strings.TrimSpace(name) == "" {
return sqlc.AgentCheckpoint{}, pcerrors.New(pcerrors.CodeUnknown, "checkpoint name is empty")
}
if runStateID.IsZero() {
return sqlc.AgentCheckpoint{}, pcerrors.New(pcerrors.CodeUnknown, "run state id is empty")
}
metaJSON := json.RawMessage(`{}`)
if meta != nil {
if b, err := json.Marshal(meta); err == nil {
metaJSON = b
}
}
return s.q.CreateAgentCheckpoint(ctx, sqlc.CreateAgentCheckpointParams{
ID: ids.New(),
ConversationID: conversationID,
Name: strings.TrimSpace(name),
RunStateID: runStateID,
MetadataJson: metaJSON,
})
}
func (s *CheckpointStore) ListCheckpoints(ctx context.Context, conversationID ids.UUID) ([]sqlc.AgentCheckpoint, error) {
if s == nil || s.q == nil {
return nil, pcerrors.New(pcerrors.CodeUnknown, "checkpoint store is not configured")
}
if conversationID.IsZero() {
return nil, pcerrors.New(pcerrors.CodeUnknown, "conversation id is empty")
}
return s.q.ListAgentCheckpointsByConversationID(ctx, sqlc.ListAgentCheckpointsByConversationIDParams{
ConversationID: conversationID,
})
}
func (s *CheckpointStore) GetCheckpoint(ctx context.Context, conversationID ids.UUID, name string) (sqlc.AgentCheckpoint, error) {
if s == nil || s.q == nil {
return sqlc.AgentCheckpoint{}, pcerrors.New(pcerrors.CodeUnknown, "checkpoint store is not configured")
}
if conversationID.IsZero() {
return sqlc.AgentCheckpoint{}, pcerrors.New(pcerrors.CodeUnknown, "conversation id is empty")
}
if strings.TrimSpace(name) == "" {
return sqlc.AgentCheckpoint{}, pcerrors.New(pcerrors.CodeUnknown, "checkpoint name is empty")
}
return s.q.GetAgentCheckpointByConversationIDAndName(ctx, sqlc.GetAgentCheckpointByConversationIDAndNameParams{
ConversationID: conversationID,
Name: strings.TrimSpace(name),
})
}

289
pkg/agent/tool_result_search.go Normal file
View file

@ -0,0 +1,289 @@
package agent
import (
"context"
"encoding/json"
"fmt"
"strings"
"charm.land/fantasy"
"github.com/sipeed/picoclaw/pkg/ids"
"github.com/sipeed/picoclaw/pkg/memory/sqlc"
"github.com/sipeed/picoclaw/pkg/pcerrors"
)
type ToolResultSearchView struct {
StartLine int `json:"start_line,omitempty" description:"Optional. 1-indexed start line (inclusive)."`
EndLine int `json:"end_line,omitempty" description:"Optional. 1-indexed end line (inclusive)."`
MaxLines int `json:"max_lines,omitempty" description:"Optional. Default 30, max 200."`
StartChunk int `json:"start_chunk,omitempty" description:"Optional. 0-indexed start chunk (inclusive)."`
EndChunk int `json:"end_chunk,omitempty" description:"Optional. 0-indexed end chunk (inclusive)."`
MaxChunks int `json:"max_chunks,omitempty" description:"Optional. Default 3, max 20."`
}
type ToolResultSearchInput struct {
ConversationID string `json:"conversation_id,omitempty" description:"Optional. Agent conversation UUID."`
RunID string `json:"run_id,omitempty" description:"Optional. Agent run UUID."`
ToolCallID string `json:"tool_call_id,omitempty" description:"Optional. Tool call id to fetch (requires run_id)."`
ToolName string `json:"tool_name,omitempty" description:"Optional. Filter by tool name."`
Query string `json:"query,omitempty" description:"Optional. Case-insensitive substring match on tool_name/tool_call_id/summary."`
Limit int `json:"limit,omitempty" description:"Optional. Default 5, max 50."`
View *ToolResultSearchView `json:"view,omitempty" description:"Optional. File view range for each result."`
}
// NewToolResultSearchTool creates the tool_result_search agent tool for
// querying previously offloaded tool results.
func NewToolResultSearchTool(q *sqlc.Queries, kv KVDelegate) fantasy.AgentTool {
return fantasy.NewAgentTool(
"tool_result_search",
"Search previously stored tool results for this agent. Supports viewing a line range or chunk range from stored results.",
func(ctx context.Context, input ToolResultSearchInput, call fantasy.ToolCall) (fantasy.ToolResponse, error) {
_ = call
if q == nil {
return fantasy.NewTextErrorResponse("db is not configured"), nil
}
if kv == nil {
return fantasy.NewTextErrorResponse("KV delegate is not configured"), nil
}
limit := input.Limit
if limit <= 0 {
limit = 5
}
if limit > 50 {
limit = 50
}
rows, err := loadToolResultRows(ctx, q, input)
if err != nil {
return fantasy.NewTextErrorResponse(err.Error()), nil
}
query := strings.ToLower(strings.TrimSpace(input.Query))
toolName := strings.TrimSpace(input.ToolName)
filtered := make([]sqlc.AgentToolResult, 0, len(rows))
for _, r := range rows {
if toolName != "" && r.ToolName != toolName {
continue
}
if query != "" {
if !strings.Contains(strings.ToLower(r.ToolName), query) &&
!strings.Contains(strings.ToLower(r.ToolCallID), query) &&
(r.Preview == nil || !strings.Contains(strings.ToLower(*r.Preview), query)) {
continue
}
}
filtered = append(filtered, r)
if len(filtered) >= limit {
break
}
}
view := input.View
startLine, endLine := normalizeLineView(view)
startChunk, endChunk := normalizeChunkView(view)
type item struct {
ID string `json:"id"`
RunID string `json:"run_id"`
StepIndex int64 `json:"step_index"`
ToolCallID string `json:"tool_call_id"`
ToolName string `json:"tool_name"`
Preview *string `json:"preview,omitempty"`
FullKey string `json:"full_key"`
ChunkCount int64 `json:"chunk_count"`
View string `json:"view"`
ViewRange map[string]int `json:"view_range"`
Metadata json.RawMessage `json:"metadata_json"`
}
out := struct {
Total int `json:"total"`
Items []item `json:"items"`
}{
Total: len(filtered),
Items: make([]item, 0, len(filtered)),
}
for _, r := range filtered {
sel, viewRange, loadErr := loadView(ctx, kv, r, startLine, endLine, startChunk, endChunk)
if loadErr != nil {
sel = "ERROR: " + loadErr.Error()
viewRange = map[string]int{
"start_line": startLine,
"end_line": endLine,
"start_chunk": startChunk,
"end_chunk": endChunk,
}
}
out.Items = append(out.Items, item{
ID: r.ID.String(),
RunID: r.RunID.String(),
StepIndex: r.StepIndex,
ToolCallID: r.ToolCallID,
ToolName: r.ToolName,
Preview: r.Preview,
FullKey: r.FullKey,
ChunkCount: r.ChunkCount,
View: sel,
ViewRange: viewRange,
Metadata: r.MetadataJson,
})
}
b, _ := json.Marshal(out)
return fantasy.NewTextResponse(string(b)), nil
},
)
}
func loadToolResultRows(ctx context.Context, q *sqlc.Queries, input ToolResultSearchInput) ([]sqlc.AgentToolResult, error) {
if q == nil {
return nil, pcerrors.New(pcerrors.CodeUnknown, "db is not configured")
}
if strings.TrimSpace(input.RunID) != "" && strings.TrimSpace(input.ToolCallID) != "" {
runID, err := ids.Parse(strings.TrimSpace(input.RunID))
if err != nil {
return nil, err
}
row, err := q.GetAgentToolResultByRunIDAndToolCallID(ctx, sqlc.GetAgentToolResultByRunIDAndToolCallIDParams{
RunID: runID,
ToolCallID: strings.TrimSpace(input.ToolCallID),
})
if err != nil {
return nil, err
}
return []sqlc.AgentToolResult{row}, nil
}
if strings.TrimSpace(input.RunID) != "" {
runID, err := ids.Parse(strings.TrimSpace(input.RunID))
if err != nil {
return nil, err
}
return q.ListAgentToolResultsByRunID(ctx, sqlc.ListAgentToolResultsByRunIDParams{
RunID: runID,
})
}
if strings.TrimSpace(input.ConversationID) != "" {
conversationID, err := ids.Parse(strings.TrimSpace(input.ConversationID))
if err != nil {
return nil, err
}
return q.ListAgentToolResultsByConversationID(ctx, sqlc.ListAgentToolResultsByConversationIDParams{
ConversationID: conversationID,
})
}
return nil, pcerrors.New(pcerrors.CodeUnknown, "conversation_id or run_id is required (and tool_call_id requires run_id)")
}
func normalizeLineView(v *ToolResultSearchView) (startLine int, endLine int) {
startLine = 1
maxLines := 30
if v == nil {
return 1, 30
}
if v.StartLine > 0 {
startLine = v.StartLine
}
if v.MaxLines > 0 {
maxLines = v.MaxLines
}
if maxLines > 200 {
maxLines = 200
}
if v.EndLine > 0 {
endLine = v.EndLine
} else {
endLine = startLine + maxLines - 1
}
if endLine < startLine {
endLine = startLine
}
return startLine, endLine
}
func normalizeChunkView(v *ToolResultSearchView) (startChunk int, endChunk int) {
startChunk = 0
maxChunks := 3
if v == nil {
return 0, 2
}
if v.StartChunk > 0 {
startChunk = v.StartChunk
}
if v.MaxChunks > 0 {
maxChunks = v.MaxChunks
}
if maxChunks > 20 {
maxChunks = 20
}
if v.EndChunk > 0 {
endChunk = v.EndChunk
} else {
endChunk = startChunk + maxChunks - 1
}
if endChunk < startChunk {
endChunk = startChunk
}
return startChunk, endChunk
}
func loadView(ctx context.Context, kv KVDelegate, row sqlc.AgentToolResult, startLine, endLine, startChunk, endChunk int) (string, map[string]int, error) {
if kv == nil {
return "", nil, pcerrors.New(pcerrors.CodeUnknown, "KV delegate is nil")
}
if row.ChunkCount > 0 {
if startChunk < 0 {
startChunk = 0
}
if int64(endChunk) >= row.ChunkCount {
endChunk = int(row.ChunkCount - 1)
}
if endChunk < startChunk {
endChunk = startChunk
}
baseDir := strings.TrimSuffix(row.FullKey, "/full.json")
var b strings.Builder
for i := startChunk; i <= endChunk; i++ {
chunkKey := fmt.Sprintf("%s/chunks/%06d.txt", baseDir, i)
part, err := kv.Get(ctx, chunkKey)
if err != nil {
return "", nil, err
}
b.Write(part)
}
return b.String(), map[string]int{"start_chunk": startChunk, "end_chunk": endChunk}, nil
}
raw, err := kv.Get(ctx, row.FullKey)
if err != nil {
return "", nil, err
}
lines := strings.Split(string(raw), "\n")
sl := startLine
el := endLine
if sl < 1 {
sl = 1
}
if el > len(lines) {
el = len(lines)
}
if el < sl {
el = sl
}
return strings.Join(lines[sl-1:el], "\n"), map[string]int{"start_line": sl, "end_line": el}, nil
}

54
pkg/pcerrors/cli.go Normal file
View file

@ -0,0 +1,54 @@
package pcerrors
import (
"fmt"
"io"
"os"
)
// CLIHandler is the PicoClaw error lifecycle boundary for the CLI.
//
// It is intentionally small: render a user-facing message and return an exit code.
// TODO: More advanced behaviors (structured logging, debug traces, redaction) can be layered on later.
type CLIHandler struct {
Writer io.Writer
}
func DefaultCLIHandler() CLIHandler {
return CLIHandler{Writer: os.Stderr}
}
func (h CLIHandler) Handle(err error) int {
if err == nil {
return 0
}
_, _ = fmt.Fprintln(h.Writer, UserMessage(err))
return ExitCode(err)
}
// UserMessage returns a friendly, stable message for humans.
//
// For structured errors, we prefer the top-level message (not the fully formatted builder.Error()).
func UserMessage(err error) string {
if eb, ok := AsErrBuilder(err); ok && eb != nil {
if eb.Msg != "" {
return eb.Msg
}
}
return err.Error()
}
// ExitCode maps error codes to process exit codes.
func ExitCode(err error) int {
if err == nil {
return 0
}
switch CodeOf(err) {
case CodeInvalidArgument, CodeFailedPrecondition, CodeOutOfRange:
return 2
case CodeUnauthenticated, CodePermissionDenied:
return 3
default:
return 1
}
}

174
pkg/pcerrors/pcerrors.go Normal file
View file

@ -0,0 +1,174 @@
package pcerrors
import (
"context"
"errors"
"fmt"
assert "github.com/ZanzyTHEbar/assert-lib"
errbuilder "github.com/ZanzyTHEbar/errbuilder-go"
)
// Code is the canonical PicoClaw error code type.
//
// We intentionally re-export errbuilder's gRPC-inspired code set so callers can
// classify errors without inventing ad-hoc sentinels.
type Code = errbuilder.ErrCode
const (
CodeCanceled = errbuilder.CodeCanceled
CodeUnknown = errbuilder.CodeUnknown
CodeInvalidArgument = errbuilder.CodeInvalidArgument
CodeDeadlineExceeded = errbuilder.CodeDeadlineExceeded
CodeNotFound = errbuilder.CodeNotFound
CodeAlreadyExists = errbuilder.CodeAlreadyExists
CodePermissionDenied = errbuilder.CodePermissionDenied
CodeResourceExhausted = errbuilder.CodeResourceExhausted
CodeFailedPrecondition = errbuilder.CodeFailedPrecondition
CodeAborted = errbuilder.CodeAborted
CodeOutOfRange = errbuilder.CodeOutOfRange
CodeUnimplemented = errbuilder.CodeUnimplemented
CodeInternal = errbuilder.CodeInternal
CodeUnavailable = errbuilder.CodeUnavailable
CodeDataLoss = errbuilder.CodeDataLoss
CodeUnauthenticated = errbuilder.CodeUnauthenticated
)
// ErrMap is a key->error bag for validation-style error details.
type ErrMap = errbuilder.ErrorMap
// Option configures an ErrBuilder before it is returned as an error.
type Option func(*buildOptions)
type buildOptions struct {
label string
cause error
details ErrMap
}
func WithLabel(label string) Option {
return func(o *buildOptions) { o.label = label }
}
func WithCause(err error) Option {
return func(o *buildOptions) { o.cause = err }
}
// WithDetail sets a single key/value detail.
//
// msg must be a string or error; other types panic (this matches errbuilder.ErrorMap.Set).
func WithDetail(key string, msg any) Option {
return func(o *buildOptions) {
o.details.Set(key, msg)
}
}
// WithDetails merges an entire error map into the error details.
func WithDetails(m ErrMap) Option {
return func(o *buildOptions) {
if m == nil {
return
}
if o.details == nil {
o.details = make(ErrMap, len(m))
}
for k, v := range m {
o.details[k] = v
}
}
}
// New constructs a structured PicoClaw error.
//
// This returns an *errbuilder.ErrBuilder which:
// - implements error
// - supports Unwrap() so stdlib errors.Is/errors.As continue to work
// - carries a Code for classification.
func New(code Code, msg string, opts ...Option) error {
o := buildOptions{}
for _, opt := range opts {
opt(&o)
}
b := errbuilder.New().WithCode(code).WithMsg(msg)
if o.label != "" {
b = b.WithLabel(o.label)
}
if o.cause != nil {
// Ensure context cancellation/deadline get codes if they weren't wrapped already.
b = b.WithCause(errbuilder.WrapIfContextError(o.cause))
}
if o.details != nil {
b = b.WithDetails(errbuilder.NewErrDetails(o.details))
}
return b
}
// Newf is like New but formats the message via fmt.Sprintf.
//
// Callers use this instead of fmt.Errorf when they want formatting
// without creating a second wrapping error layer.
func Newf(code Code, format string, args ...any) error {
return New(code, fmt.Sprintf(format, args...))
}
// Wrap wraps err with a new structured error code and message.
func Wrap(code Code, err error, msg string, opts ...Option) error {
if err == nil {
return nil
}
return New(code, msg, append(opts, WithCause(err))...)
}
// Wrapf wraps err with a formatted message.
func Wrapf(code Code, err error, format string, args ...any) error {
return Wrap(code, err, fmt.Sprintf(format, args...))
}
// CodeOf returns the structured code for err if it is (or wraps) an ErrBuilder.
// Otherwise it returns CodeUnknown.
func CodeOf(err error) Code {
return errbuilder.CodeOf(err)
}
// Is is a thin wrapper over errors.Is.
func Is(err error, target error) bool {
return errors.Is(err, target)
}
// As is a type-safe wrapper over errors.As.
//
// Go's errors.As will panic if the provided target is not a pointer to an
// interface or a type implementing error.
//
// By constraining T to error, we eliminate the most common foot-gun at compile time.
func As[T error](err error) (T, bool) {
var target T
if err == nil {
return target, false
}
if errors.As(err, &target) {
return target, true
}
return target, false
}
// AsInto mirrors errors.As but keeps the type safety of T error.
func AsInto[T error](err error, target *T) bool {
if err == nil || target == nil {
return false
}
return errors.As(err, target)
}
// AsErrBuilder extracts an underlying *errbuilder.ErrBuilder if present.
func AsErrBuilder(err error) (*errbuilder.ErrBuilder, bool) {
return As[*errbuilder.ErrBuilder](err)
}
// Assert integrates assert-lib so callers can opt into lightweight runtime checks.
//
// We default to production-friendly formatting, and we keep the library's safe-by-default behavior.
func Assert(ctx context.Context, condition bool, msg string, opts ...assert.Option) {
assert.Assert(ctx, condition, msg, append([]assert.Option{assert.WithProductionDefaults()}, opts...)...)
}

259
pkg/security/jsonextract.go Normal file
View file

@ -0,0 +1,259 @@
// Package security provides hardened input handling for LLM-generated content.
package security
import (
"encoding/json"
"errors"
"fmt"
"strings"
)
// ErrNoJSON is returned when no valid JSON object is found in the input.
var ErrNoJSON = errors.New("no valid JSON object found in text")
// ErrInputTooLarge is returned when the input exceeds the maximum allowed size.
var ErrInputTooLarge = errors.New("input exceeds maximum allowed size")
const (
defaultMaxInputBytes = 64 * 1024 // 64KB
)
// ExtractJSONOptions configures JSON extraction behavior.
type ExtractJSONOptions struct {
MaxInputBytes int // Maximum input size in bytes. Default: 64KB.
DisallowUnknownFields bool // Reject JSON with keys not in the target struct.
}
// ExtractJSON extracts the first valid JSON object from LLM-generated text and
// unmarshals it into dest. It handles common LLM output patterns:
// - Raw JSON
// - JSON wrapped in ```json ... ``` code fences
// - JSON embedded in prose text
//
// It applies size limits and optionally rejects unknown fields, defending
// against prompt injection attacks that attempt to smuggle extra keys.
func ExtractJSON(text string, dest interface{}, opts *ExtractJSONOptions) error {
if opts == nil {
opts = &ExtractJSONOptions{}
}
maxBytes := opts.MaxInputBytes
if maxBytes <= 0 {
maxBytes = defaultMaxInputBytes
}
if len(text) > maxBytes {
return fmt.Errorf("%w: %d bytes (max %d)", ErrInputTooLarge, len(text), maxBytes)
}
cleaned := extractJSONString(text)
if cleaned == "" {
return ErrNoJSON
}
dec := json.NewDecoder(strings.NewReader(cleaned))
if opts.DisallowUnknownFields {
dec.DisallowUnknownFields()
}
if err := dec.Decode(dest); err != nil {
return fmt.Errorf("json decode: %w", err)
}
return nil
}
// extractJSONString isolates the JSON object from surrounding text.
// Strategy (in priority order):
// 1. Try the text as-is (after trim)
// 2. Extract from ```json ... ``` code fence
// 3. Find the first { ... } balanced brace pair
func extractJSONString(text string) string {
text = strings.TrimSpace(text)
if isJSON(text) {
return text
}
if extracted := extractFromCodeFence(text); extracted != "" && isJSON(extracted) {
return extracted
}
if extracted := extractFirstBraced(text); extracted != "" && isJSON(extracted) {
return extracted
}
return ""
}
func isJSON(s string) bool {
s = strings.TrimSpace(s)
if len(s) < 2 {
return false
}
return (s[0] == '{' && s[len(s)-1] == '}') || (s[0] == '[' && s[len(s)-1] == ']')
}
// extractFromCodeFence extracts content from the first ```json ... ``` fence.
func extractFromCodeFence(text string) string {
lower := strings.ToLower(text)
markers := []string{"```json", "```"}
for _, marker := range markers {
idx := strings.Index(lower, marker)
if idx < 0 {
continue
}
start := idx + len(marker)
rest := text[start:]
endIdx := strings.Index(rest, "```")
if endIdx < 0 {
continue
}
return strings.TrimSpace(rest[:endIdx])
}
return ""
}
// extractFirstBraced finds the first balanced { ... } in text.
// Handles nested braces and string literals containing braces.
func extractFirstBraced(text string) string {
start := strings.IndexByte(text, '{')
if start < 0 {
return ""
}
depth := 0
inString := false
escaped := false
for i := start; i < len(text); i++ {
ch := text[i]
if escaped {
escaped = false
continue
}
if ch == '\\' && inString {
escaped = true
continue
}
if ch == '"' {
inString = !inString
continue
}
if inString {
continue
}
switch ch {
case '{':
depth++
case '}':
depth--
if depth == 0 {
return text[start : i+1]
}
}
}
return ""
}
// SanitizeToolArgs validates that tool arguments conform to expected types and
// constraints. It returns the sanitized arguments or an error.
func SanitizeToolArgs(args map[string]interface{}, schema map[string]ArgSpec) (map[string]interface{}, error) {
sanitized := make(map[string]interface{}, len(schema))
for name, spec := range schema {
val, exists := args[name]
if !exists || val == nil {
if spec.Required {
return nil, fmt.Errorf("missing required argument: %s", name)
}
if spec.Default != nil {
sanitized[name] = spec.Default
}
continue
}
coerced, err := coerceArg(val, spec.Type)
if err != nil {
return nil, fmt.Errorf("argument %q: %w", name, err)
}
if spec.MaxLength > 0 {
if s, ok := coerced.(string); ok && len(s) > spec.MaxLength {
return nil, fmt.Errorf("argument %q exceeds max length %d", name, spec.MaxLength)
}
}
sanitized[name] = coerced
}
return sanitized, nil
}
// ArgSpec defines validation constraints for a single tool argument.
type ArgSpec struct {
Type ArgType // Expected type.
Required bool // Whether the argument is required.
MaxLength int // Maximum string length (0 = unlimited).
Default interface{} // Default value if not provided.
}
// ArgType represents the expected type of a tool argument.
type ArgType int
const (
ArgString ArgType = iota
ArgInt
ArgFloat
ArgBool
ArgObject
ArgArray
)
func coerceArg(val interface{}, expected ArgType) (interface{}, error) {
switch expected {
case ArgString:
switch v := val.(type) {
case string:
return v, nil
case float64:
return fmt.Sprintf("%g", v), nil
default:
return nil, fmt.Errorf("expected string, got %T", val)
}
case ArgInt:
switch v := val.(type) {
case float64:
return int(v), nil
case int:
return v, nil
default:
return nil, fmt.Errorf("expected integer, got %T", val)
}
case ArgFloat:
switch v := val.(type) {
case float64:
return v, nil
case int:
return float64(v), nil
default:
return nil, fmt.Errorf("expected number, got %T", val)
}
case ArgBool:
if b, ok := val.(bool); ok {
return b, nil
}
return nil, fmt.Errorf("expected boolean, got %T", val)
case ArgObject:
if m, ok := val.(map[string]interface{}); ok {
return m, nil
}
return nil, fmt.Errorf("expected object, got %T", val)
case ArgArray:
if a, ok := val.([]interface{}); ok {
return a, nil
}
return nil, fmt.Errorf("expected array, got %T", val)
default:
return val, nil
}
}

232
pkg/security/jsonextract_test.go Normal file
View file

@ -0,0 +1,232 @@
package security
import (
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestExtractJSON_RawJSON(t *testing.T) {
tests := []struct {
name string
input string
wantKey string
wantVal interface{}
}{
{"clean object", `{"importance": 0.8}`, "importance", 0.8},
{"with whitespace", ` {"key": "val"} `, "key", "val"},
{"nested", `{"outer": {"inner": true}}`, "outer", map[string]interface{}{"inner": true}},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
var result map[string]interface{}
err := ExtractJSON(tc.input, &result, nil)
require.NoError(t, err)
assert.Equal(t, tc.wantVal, result[tc.wantKey])
})
}
}
func TestExtractJSON_CodeFence(t *testing.T) {
tests := []struct {
name string
input string
}{
{
"json fence",
"Here is the result:\n```json\n{\"score\": 42}\n```\n",
},
{
"plain fence",
"```\n{\"score\": 42}\n```",
},
{
"fence with prose before and after",
"I analyzed the data.\n```json\n{\"score\": 42}\n```\nDone!",
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
var result map[string]interface{}
err := ExtractJSON(tc.input, &result, nil)
require.NoError(t, err)
assert.Equal(t, float64(42), result["score"])
})
}
}
func TestExtractJSON_EmbeddedInProse(t *testing.T) {
input := `Based on my analysis, the result is {"importance": 0.9, "sector": "semantic"} which indicates high relevance.`
var result struct {
Importance float64 `json:"importance"`
Sector string `json:"sector"`
}
err := ExtractJSON(input, &result, nil)
require.NoError(t, err)
assert.Equal(t, 0.9, result.Importance)
assert.Equal(t, "semantic", result.Sector)
}
func TestExtractJSON_NestedBracesInStrings(t *testing.T) {
input := `{"content": "function() { return {}; }", "count": 1}`
var result map[string]interface{}
err := ExtractJSON(input, &result, nil)
require.NoError(t, err)
assert.Equal(t, "function() { return {}; }", result["content"])
assert.Equal(t, float64(1), result["count"])
}
func TestExtractJSON_InjectionAttempts(t *testing.T) {
tests := []struct {
name string
input string
opts *ExtractJSONOptions
check func(t *testing.T, err error)
}{
{
"oversized input",
strings.Repeat("x", 100*1024),
nil,
func(t *testing.T, err error) {
assert.ErrorIs(t, err, ErrInputTooLarge)
},
},
{
"no json at all",
"This is just prose with no JSON.",
nil,
func(t *testing.T, err error) {
assert.ErrorIs(t, err, ErrNoJSON)
},
},
{
"unknown fields rejected",
`{"importance": 0.5, "injected_key": "malicious"}`,
&ExtractJSONOptions{DisallowUnknownFields: true},
func(t *testing.T, err error) {
assert.Error(t, err)
assert.Contains(t, err.Error(), "unknown field")
},
},
{
"custom size limit",
`{"key": "val"}`,
&ExtractJSONOptions{MaxInputBytes: 5},
func(t *testing.T, err error) {
assert.ErrorIs(t, err, ErrInputTooLarge)
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
var result struct {
Importance float64 `json:"importance"`
}
err := ExtractJSON(tc.input, &result, tc.opts)
tc.check(t, err)
})
}
}
func TestExtractJSON_EmptyInput(t *testing.T) {
var result map[string]interface{}
err := ExtractJSON("", &result, nil)
assert.ErrorIs(t, err, ErrNoJSON)
}
func TestExtractJSON_MultipleFences_TakesFirst(t *testing.T) {
input := "```json\n{\"first\": true}\n```\nmore text\n```json\n{\"second\": true}\n```"
var result map[string]interface{}
err := ExtractJSON(input, &result, nil)
require.NoError(t, err)
assert.Equal(t, true, result["first"])
_, hasSecond := result["second"]
assert.False(t, hasSecond)
}
func TestSanitizeToolArgs_ValidInput(t *testing.T) {
schema := map[string]ArgSpec{
"path": {Type: ArgString, Required: true, MaxLength: 256},
"content": {Type: ArgString, Required: true},
"mode": {Type: ArgString, Required: false, Default: "overwrite"},
}
args := map[string]interface{}{
"path": "/tmp/test.txt",
"content": "hello world",
}
result, err := SanitizeToolArgs(args, schema)
require.NoError(t, err)
assert.Equal(t, "/tmp/test.txt", result["path"])
assert.Equal(t, "hello world", result["content"])
assert.Equal(t, "overwrite", result["mode"])
}
func TestSanitizeToolArgs_MissingRequired(t *testing.T) {
schema := map[string]ArgSpec{
"path": {Type: ArgString, Required: true},
}
_, err := SanitizeToolArgs(map[string]interface{}{}, schema)
assert.Error(t, err)
assert.Contains(t, err.Error(), "missing required argument")
}
func TestSanitizeToolArgs_ExceedsMaxLength(t *testing.T) {
schema := map[string]ArgSpec{
"cmd": {Type: ArgString, Required: true, MaxLength: 10},
}
_, err := SanitizeToolArgs(map[string]interface{}{"cmd": "very long command string"}, schema)
assert.Error(t, err)
assert.Contains(t, err.Error(), "max length")
}
func TestSanitizeToolArgs_TypeCoercion(t *testing.T) {
tests := []struct {
name string
argType ArgType
input interface{}
expected interface{}
wantErr bool
}{
{"string from string", ArgString, "hello", "hello", false},
{"string from float", ArgString, float64(42), "42", false},
{"int from float", ArgInt, float64(42), 42, false},
{"float from int", ArgFloat, 42, float64(42), false},
{"bool valid", ArgBool, true, true, false},
{"bool invalid", ArgBool, "true", nil, true},
{"object valid", ArgObject, map[string]interface{}{"k": "v"}, map[string]interface{}{"k": "v"}, false},
{"array valid", ArgArray, []interface{}{"a"}, []interface{}{"a"}, false},
{"array invalid", ArgArray, "not array", nil, true},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
schema := map[string]ArgSpec{
"arg": {Type: tc.argType, Required: true},
}
result, err := SanitizeToolArgs(map[string]interface{}{"arg": tc.input}, schema)
if tc.wantErr {
assert.Error(t, err)
} else {
require.NoError(t, err)
assert.Equal(t, tc.expected, result["arg"])
}
})
}
}
func TestExtractFirstBraced_EscapedQuotes(t *testing.T) {
input := `{"msg": "say \"hello\" world"}`
result := extractFirstBraced(input)
assert.Equal(t, input, result)
}
func TestExtractFirstBraced_UnbalancedBraces(t *testing.T) {
input := `text { not closed`
result := extractFirstBraced(input)
assert.Empty(t, result)
}

123
pkg/security/redact.go Normal file
View file

@ -0,0 +1,123 @@
package security
import (
"regexp"
"strings"
)
// Redactor strips sensitive patterns from text before it reaches logs or storage.
type Redactor struct {
patterns []*redactPattern
}
type redactPattern struct {
re *regexp.Regexp
label string
}
// NewRedactor builds a Redactor with the default set of PII/secret patterns.
func NewRedactor() *Redactor {
return &Redactor{patterns: defaultPatterns()}
}
// Redact replaces all matched secrets/PII in text with [REDACTED:<label>].
func (r *Redactor) Redact(text string) string {
for _, p := range r.patterns {
text = p.re.ReplaceAllString(text, "[REDACTED:"+p.label+"]")
}
return text
}
// ContainsSensitive returns true if any pattern matches the text.
func (r *Redactor) ContainsSensitive(text string) bool {
for _, p := range r.patterns {
if p.re.MatchString(text) {
return true
}
}
return false
}
func defaultPatterns() []*redactPattern {
defs := []struct {
pattern string
label string
}{
// Anthropic API keys (must come before generic sk- pattern)
{`sk-ant-[A-Za-z0-9_-]{20,}`, "ANTHROPIC_KEY"},
// OpenAI API keys (sk-..., sk-proj-...)
{`sk-[A-Za-z0-9_-]{20,}`, "OPENAI_KEY"},
// Generic Bearer tokens in header-like context
{`Bearer\s+[A-Za-z0-9._~+/=-]{20,}`, "BEARER_TOKEN"},
// AWS access key IDs (AKIA...)
{`AKIA[0-9A-Z]{16}`, "AWS_ACCESS_KEY"},
// AWS secret keys (40 chars base64-ish after common prefixes)
{`(?i)aws_secret_access_key\s*[=:]\s*[A-Za-z0-9/+=]{40}`, "AWS_SECRET_KEY"},
// GitHub tokens (ghp_, gho_, ghu_, ghs_, ghr_) — must come before generic SECRET_VALUE
{`gh[pousr]_[A-Za-z0-9_]{36,}`, "GITHUB_TOKEN"},
// Slack tokens (xoxb-, xoxp-, xoxs-, xoxa-)
{`xox[bpsa]-[A-Za-z0-9-]{10,}`, "SLACK_TOKEN"},
// Generic secret/password in key=value or key: value lines
{`(?i)(password|secret|token|api_key|apikey|api-key)\s*[=:]\s*\S{8,}`, "SECRET_VALUE"},
// SSH private key markers
{`-----BEGIN (RSA |EC |OPENSSH |DSA )?PRIVATE KEY-----`, "SSH_PRIVATE_KEY"},
// JWT tokens (three base64url segments separated by dots)
{`eyJ[A-Za-z0-9_-]{10,}\.eyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}`, "JWT"},
// Email addresses (basic)
{`[a-zA-Z0-9._%+\-]+@[a-zA-Z0-9.\-]+\.[a-zA-Z]{2,}`, "EMAIL"},
// Credit card numbers (13-19 digits with optional separators)
{`\b(?:\d[ -]*?){13,19}\b`, "CREDIT_CARD"},
// US Social Security Numbers (XXX-XX-XXXX)
{`\b\d{3}-\d{2}-\d{4}\b`, "SSN"},
// IP addresses (to catch accidental logging of internal IPs)
// Only match private ranges to avoid over-redacting
{`\b(?:10\.\d{1,3}\.\d{1,3}\.\d{1,3}|172\.(?:1[6-9]|2\d|3[01])\.\d{1,3}\.\d{1,3}|192\.168\.\d{1,3}\.\d{1,3})\b`, "PRIVATE_IP"},
}
patterns := make([]*redactPattern, 0, len(defs))
for _, d := range defs {
re, err := regexp.Compile(d.pattern)
if err != nil {
continue
}
patterns = append(patterns, &redactPattern{re: re, label: d.label})
}
return patterns
}
// RedactMap redacts all values in a string map, returning a new map.
func (r *Redactor) RedactMap(m map[string]interface{}) map[string]interface{} {
out := make(map[string]interface{}, len(m))
for k, v := range m {
switch val := v.(type) {
case string:
out[k] = r.Redact(val)
case map[string]interface{}:
out[k] = r.RedactMap(val)
default:
out[k] = v
}
}
return out
}
// MaskKey partially masks a key, showing only a prefix for identification.
func MaskKey(key string) string {
if len(key) <= 8 {
return strings.Repeat("*", len(key))
}
return key[:8] + strings.Repeat("*", len(key)-8)
}

126
pkg/security/redact_test.go Normal file
View file

@ -0,0 +1,126 @@
package security
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestRedactor_APIKeys(t *testing.T) {
r := NewRedactor()
tests := []struct {
name string
input string
label string
}{
{"OpenAI key", "my key is sk-abc123XYZdefghijklmnopqrs", "OPENAI_KEY"},
{"OpenAI proj key", "using sk-proj-abc123XYZdefghijklmnopqrs", "OPENAI_KEY"},
{"Anthropic key", "key: sk-ant-abc123XYZdefghijklmnopqrs", "ANTHROPIC_KEY"},
{"GitHub token", "using ghp_abcdefghijklmnopqrstuvwxyz1234567890 here", "GITHUB_TOKEN"},
{"Slack bot token", "token xoxb-123456789-abcdefghijk", "SLACK_TOKEN"},
{"AWS access key", "AKIAIOSFODNN7EXAMPLE", "AWS_ACCESS_KEY"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
out := r.Redact(tc.input)
assert.Contains(t, out, "[REDACTED:"+tc.label+"]")
assert.True(t, r.ContainsSensitive(tc.input))
})
}
}
func TestRedactor_Bearer(t *testing.T) {
r := NewRedactor()
input := "Authorization: Bearer eyJhbGciOiJIUzI1NiJ9.xxxxx"
out := r.Redact(input)
assert.Contains(t, out, "[REDACTED:")
}
func TestRedactor_JWT(t *testing.T) {
r := NewRedactor()
jwt := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.dozjgNryP4J3jVmNHl0w5N_XgL0n3I9PlFUP0THsR8U"
out := r.Redact(jwt)
assert.Contains(t, out, "[REDACTED:JWT]")
}
func TestRedactor_PII(t *testing.T) {
r := NewRedactor()
tests := []struct {
name string
input string
label string
}{
{"email", "contact john.doe@example.com please", "EMAIL"},
{"SSN", "SSN: 123-45-6789", "SSN"},
{"private IP", "connecting to 192.168.1.100", "PRIVATE_IP"},
{"10.x IP", "server at 10.0.0.1", "PRIVATE_IP"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
out := r.Redact(tc.input)
assert.Contains(t, out, "[REDACTED:"+tc.label+"]")
})
}
}
func TestRedactor_SecretValues(t *testing.T) {
r := NewRedactor()
tests := []struct {
name string
input string
}{
{"password=", "password=SuperSecret123!"},
{"api_key:", "api_key: sk_live_1234567890"},
{"token=", "token=abcdefgh12345678"},
{"SSH key header", "-----BEGIN RSA PRIVATE KEY-----"},
{"SSH key EC", "-----BEGIN EC PRIVATE KEY-----"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
out := r.Redact(tc.input)
assert.Contains(t, out, "[REDACTED:")
})
}
}
func TestRedactor_SafeText(t *testing.T) {
r := NewRedactor()
safe := "This is a normal log message about processing 42 items."
assert.Equal(t, safe, r.Redact(safe))
assert.False(t, r.ContainsSensitive(safe))
}
func TestRedactor_RedactMap(t *testing.T) {
r := NewRedactor()
m := map[string]interface{}{
"command": "curl -H 'Authorization: Bearer eyJhbGciOiJIUzI1NiJ9.xxxxxxxxxxxxxxxxxxxx.yyyyyyyy'",
"output": "normal output",
"nested": map[string]interface{}{
"secret": "password=hunter2hunter2",
},
"count": 42,
}
out := r.RedactMap(m)
assert.Contains(t, out["command"].(string), "[REDACTED:")
assert.Equal(t, "normal output", out["output"])
nested := out["nested"].(map[string]interface{})
assert.Contains(t, nested["secret"].(string), "[REDACTED:")
assert.Equal(t, 42, out["count"])
}
func TestMaskKey(t *testing.T) {
tests := []struct {
input string
want string
}{
{"sk-abc123XYZdefghijklmno", "sk-abc12****************"},
{"short", "*****"},
{"12345678", "********"},
{"123456789", "12345678*"},
}
for _, tc := range tests {
t.Run(tc.input, func(t *testing.T) {
assert.Equal(t, tc.want, MaskKey(tc.input))
})
}
}

116
pkg/security/urlguard.go Normal file
View file

@ -0,0 +1,116 @@
package security
import (
"fmt"
"net"
"net/url"
"strings"
)
// ErrBlockedURL is returned when a URL targets a blocked network.
var ErrBlockedURL = fmt.Errorf("URL targets a blocked network")
// ValidateURL checks that a URL is safe to fetch, blocking internal/private
// IPs, cloud metadata endpoints, and loopback addresses that could be used
// for SSRF attacks.
func ValidateURL(rawURL string) error {
parsed, err := url.Parse(rawURL)
if err != nil {
return fmt.Errorf("invalid URL: %w", err)
}
scheme := strings.ToLower(parsed.Scheme)
if scheme != "http" && scheme != "https" {
return fmt.Errorf("%w: only http/https schemes allowed, got %q", ErrBlockedURL, scheme)
}
host := parsed.Hostname()
if host == "" {
return fmt.Errorf("%w: empty hostname", ErrBlockedURL)
}
if isBlockedHost(host) {
return fmt.Errorf("%w: host %q is blocked", ErrBlockedURL, host)
}
ips, err := net.LookupHost(host)
if err != nil {
return fmt.Errorf("%w: DNS resolution failed for %q: %v", ErrBlockedURL, host, err)
}
for _, ipStr := range ips {
ip := net.ParseIP(ipStr)
if ip == nil {
continue
}
if isBlockedIP(ip) {
return fmt.Errorf("%w: resolved IP %s is in a blocked range", ErrBlockedURL, ipStr)
}
}
return nil
}
// isBlockedHost checks hostnames that are always blocked regardless of resolution.
func isBlockedHost(host string) bool {
lower := strings.ToLower(host)
blockedHosts := []string{
"localhost",
"metadata.google.internal",
"metadata.google",
"169.254.169.254",
}
for _, blocked := range blockedHosts {
if lower == blocked {
return true
}
}
blockedSuffixes := []string{
".internal",
".local",
".localhost",
}
for _, suffix := range blockedSuffixes {
if strings.HasSuffix(lower, suffix) {
return true
}
}
return false
}
// isBlockedIP returns true if the IP is in a private, loopback, link-local,
// or otherwise blocked range.
func isBlockedIP(ip net.IP) bool {
if ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() ||
ip.IsLinkLocalMulticast() || ip.IsUnspecified() {
return true
}
blocked := []struct {
network string
cidr string
}{
{"AWS metadata", "169.254.169.254/32"},
{"CGNAT", "100.64.0.0/10"},
{"Benchmarking", "198.18.0.0/15"},
{"Documentation", "192.0.2.0/24"},
{"Documentation2", "198.51.100.0/24"},
{"Documentation3", "203.0.113.0/24"},
{"IPv6 unique local", "fc00::/7"},
}
for _, b := range blocked {
_, cidr, err := net.ParseCIDR(b.cidr)
if err != nil {
continue
}
if cidr.Contains(ip) {
return true
}
}
return false
}

121
pkg/security/urlguard_test.go Normal file
View file

@ -0,0 +1,121 @@
package security
import (
"net"
"testing"
"github.com/stretchr/testify/assert"
)
func TestValidateURL_AllowedURLs(t *testing.T) {
tests := []string{
"https://example.com",
"https://api.openai.com/v1/chat",
"https://www.google.com/robots.txt",
"https://github.com/user/repo",
}
for _, u := range tests {
t.Run(u, func(t *testing.T) {
err := ValidateURL(u)
assert.NoError(t, err)
})
}
}
func TestValidateURL_BlockedSchemes(t *testing.T) {
tests := []string{
"file:///etc/passwd",
"ftp://internal.server/data",
"gopher://evil.com/payload",
"javascript:alert(1)",
}
for _, u := range tests {
t.Run(u, func(t *testing.T) {
err := ValidateURL(u)
assert.ErrorIs(t, err, ErrBlockedURL)
})
}
}
func TestValidateURL_BlockedHosts(t *testing.T) {
tests := []struct {
name string
url string
}{
{"localhost", "http://localhost/api"},
{"metadata google", "http://metadata.google.internal/computeMetadata/v1/"},
{"AWS metadata IP", "http://169.254.169.254/latest/meta-data/"},
{"local suffix", "http://app.local/api"},
{"localhost suffix", "http://service.localhost/api"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
err := ValidateURL(tc.url)
assert.ErrorIs(t, err, ErrBlockedURL)
})
}
}
func TestValidateURL_BlockedIPs(t *testing.T) {
tests := []struct {
name string
url string
}{
{"loopback v4", "http://127.0.0.1/api"},
{"loopback v6", "http://[::1]/api"},
{"private 10.x", "http://10.0.0.1/secret"},
{"private 172.x", "http://172.16.0.1/admin"},
{"private 192.168.x", "http://192.168.1.1/config"},
{"link-local", "http://169.254.1.1/data"},
{"unspecified", "http://0.0.0.0/"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
err := ValidateURL(tc.url)
assert.ErrorIs(t, err, ErrBlockedURL)
})
}
}
func TestValidateURL_EmptyAndInvalid(t *testing.T) {
tests := []struct {
name string
url string
}{
{"empty", ""},
{"no scheme", "example.com"},
{"whitespace", " "},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
err := ValidateURL(tc.url)
assert.Error(t, err)
})
}
}
func TestIsBlockedIP(t *testing.T) {
tests := []struct {
ip string
blocked bool
}{
{"8.8.8.8", false},
{"1.1.1.1", false},
{"127.0.0.1", true},
{"10.0.0.1", true},
{"172.16.0.1", true},
{"192.168.0.1", true},
{"169.254.169.254", true},
{"100.64.0.1", true},
{"0.0.0.0", true},
}
for _, tc := range tests {
t.Run(tc.ip, func(t *testing.T) {
ip := net.ParseIP(tc.ip)
if ip == nil {
t.Fatalf("invalid IP: %s", tc.ip)
}
assert.Equal(t, tc.blocked, isBlockedIP(ip))
})
}
}

99
pkg/security/vault.go Normal file
View file

@ -0,0 +1,99 @@
package security
import (
"crypto/rand"
"encoding/base64"
"errors"
"fmt"
"io"
"golang.org/x/crypto/chacha20poly1305"
)
var (
ErrDecryptFailed = errors.New("decryption failed: ciphertext tampered or wrong key")
ErrKeyLength = fmt.Errorf("key must be exactly %d bytes", chacha20poly1305.KeySize)
)
// Vault encrypts and decrypts secrets using XChaCha20-Poly1305 (AEAD).
// XChaCha20 uses a 24-byte nonce, making random nonces safe even at high volume.
type Vault struct {
key []byte
}
// NewVault creates a Vault from a 32-byte key. Returns an error if the key
// length is wrong.
func NewVault(key []byte) (*Vault, error) {
if len(key) != chacha20poly1305.KeySize {
return nil, ErrKeyLength
}
k := make([]byte, chacha20poly1305.KeySize)
copy(k, key)
return &Vault{key: k}, nil
}
// GenerateKey produces a cryptographically random 32-byte key suitable for NewVault.
func GenerateKey() ([]byte, error) {
key := make([]byte, chacha20poly1305.KeySize)
if _, err := io.ReadFull(rand.Reader, key); err != nil {
return nil, fmt.Errorf("generate key: %w", err)
}
return key, nil
}
// Encrypt seals plaintext using XChaCha20-Poly1305 and returns a base64url-encoded
// string containing nonce + ciphertext.
func (v *Vault) Encrypt(plaintext []byte) (string, error) {
aead, err := chacha20poly1305.NewX(v.key)
if err != nil {
return "", fmt.Errorf("create cipher: %w", err)
}
nonce := make([]byte, aead.NonceSize())
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
return "", fmt.Errorf("generate nonce: %w", err)
}
ciphertext := aead.Seal(nonce, nonce, plaintext, nil)
return base64.RawURLEncoding.EncodeToString(ciphertext), nil
}
// Decrypt decodes a base64url string and opens it with XChaCha20-Poly1305.
func (v *Vault) Decrypt(encoded string) ([]byte, error) {
data, err := base64.RawURLEncoding.DecodeString(encoded)
if err != nil {
return nil, fmt.Errorf("decode base64: %w", err)
}
aead, err := chacha20poly1305.NewX(v.key)
if err != nil {
return nil, fmt.Errorf("create cipher: %w", err)
}
if len(data) < aead.NonceSize() {
return nil, ErrDecryptFailed
}
nonce := data[:aead.NonceSize()]
ciphertext := data[aead.NonceSize():]
plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
if err != nil {
return nil, ErrDecryptFailed
}
return plaintext, nil
}
// EncryptString is a convenience wrapper around Encrypt.
func (v *Vault) EncryptString(s string) (string, error) {
return v.Encrypt([]byte(s))
}
// DecryptString is a convenience wrapper around Decrypt.
func (v *Vault) DecryptString(encoded string) (string, error) {
plaintext, err := v.Decrypt(encoded)
if err != nil {
return "", err
}
return string(plaintext), nil
}

103
pkg/security/vault_test.go Normal file
View file

@ -0,0 +1,103 @@
package security
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestVault_RoundTrip(t *testing.T) {
key, err := GenerateKey()
require.NoError(t, err)
require.Len(t, key, 32)
v, err := NewVault(key)
require.NoError(t, err)
tests := []string{
"sk-abc123XYZdefghijklmnopqrs",
"",
"short",
"a longer secret with spaces and special chars: !@#$%^&*()",
}
for _, secret := range tests {
t.Run(secret, func(t *testing.T) {
enc, err := v.EncryptString(secret)
require.NoError(t, err)
assert.NotEqual(t, secret, enc)
dec, err := v.DecryptString(enc)
require.NoError(t, err)
assert.Equal(t, secret, dec)
})
}
}
func TestVault_DifferentCiphertexts(t *testing.T) {
key, _ := GenerateKey()
v, _ := NewVault(key)
enc1, _ := v.EncryptString("same secret")
enc2, _ := v.EncryptString("same secret")
assert.NotEqual(t, enc1, enc2, "random nonces should produce different ciphertexts")
}
func TestVault_WrongKey(t *testing.T) {
key1, _ := GenerateKey()
key2, _ := GenerateKey()
v1, _ := NewVault(key1)
v2, _ := NewVault(key2)
enc, err := v1.EncryptString("secret data")
require.NoError(t, err)
_, err = v2.DecryptString(enc)
assert.ErrorIs(t, err, ErrDecryptFailed)
}
func TestVault_TamperedCiphertext(t *testing.T) {
key, _ := GenerateKey()
v, _ := NewVault(key)
enc, err := v.EncryptString("tamper test")
require.NoError(t, err)
tampered := enc[:len(enc)-2] + "XX"
_, err = v.DecryptString(tampered)
assert.Error(t, err)
}
func TestVault_InvalidKeyLength(t *testing.T) {
_, err := NewVault([]byte("too-short"))
assert.ErrorIs(t, err, ErrKeyLength)
_, err = NewVault(make([]byte, 64))
assert.ErrorIs(t, err, ErrKeyLength)
}
func TestVault_EmptyInput(t *testing.T) {
key, _ := GenerateKey()
v, _ := NewVault(key)
enc, err := v.EncryptString("")
require.NoError(t, err)
dec, err := v.DecryptString(enc)
require.NoError(t, err)
assert.Equal(t, "", dec)
}
func TestVault_BinaryData(t *testing.T) {
key, _ := GenerateKey()
v, _ := NewVault(key)
binary := []byte{0x00, 0x01, 0xFF, 0xFE, 0x80, 0x7F}
enc, err := v.Encrypt(binary)
require.NoError(t, err)
dec, err := v.Decrypt(enc)
require.NoError(t, err)
assert.Equal(t, binary, dec)
}

193
pkg/worker/worker.go Normal file
View file

@ -0,0 +1,193 @@
package worker
import (
"context"
"database/sql"
"fmt"
"os"
"time"
"github.com/rs/zerolog"
"github.com/sipeed/picoclaw/pkg/memory/sqlc"
"github.com/sipeed/picoclaw/pkg/pcerrors"
)
// HandlerFunc processes a single claimed job. Returning nil marks it succeeded.
type HandlerFunc func(ctx context.Context, q *sqlc.Queries, job sqlc.Job) error
// Options configures the worker loop.
type Options struct {
Handlers map[string]HandlerFunc
LockedBy string
Sleep time.Duration
Backoff func(attempt int) time.Duration
Now func() time.Time
Logger *zerolog.Logger
}
func (o *Options) lockedBy() string {
if o != nil && o.LockedBy != "" {
return o.LockedBy
}
hn, err := os.Hostname()
if err == nil && hn != "" {
return hn
}
return "worker"
}
func (o *Options) now() time.Time {
if o != nil && o.Now != nil {
return o.Now()
}
return time.Now().UTC()
}
func (o *Options) sleep() time.Duration {
if o != nil && o.Sleep > 0 {
return o.Sleep
}
return 2 * time.Second
}
func (o *Options) backoff(attempt int) time.Duration {
if o != nil && o.Backoff != nil {
return o.Backoff(attempt)
}
d := time.Duration(1<<min(attempt-1, 5)) * time.Second
if d > 5*time.Minute {
return 5 * time.Minute
}
return d
}
func (o *Options) logger() zerolog.Logger {
if o != nil && o.Logger != nil {
return *o.Logger
}
return zerolog.Nop()
}
// RunOnce claims and executes a single job if available.
func RunOnce(ctx context.Context, q *sqlc.Queries, opts *Options) error {
if q == nil {
return pcerrors.New(pcerrors.CodeFailedPrecondition, "queries is nil")
}
l := opts.logger()
lockedBy := opts.lockedBy()
handlers := map[string]HandlerFunc{}
if opts != nil && opts.Handlers != nil {
handlers = opts.Handlers
}
jobID, err := q.FindNextRunnableJob(ctx)
if err != nil {
if err == sql.ErrNoRows {
return nil
}
return err
}
job, err := q.ClaimJobByID(ctx, sqlc.ClaimJobByIDParams{
ID: jobID,
LockedBy: &lockedBy,
})
if err != nil {
if err == sql.ErrNoRows {
return nil
}
return err
}
handler := handlers[job.Kind]
if handler == nil {
msg := fmt.Sprintf("no handler for kind=%s", job.Kind)
l.Warn().Str("job_id", job.ID.String()).Msg(msg)
_, _ = q.MarkJobFailed(ctx, sqlc.MarkJobFailedParams{
LastError: &msg,
ID: job.ID,
})
return nil
}
l.Info().
Str("job_id", job.ID.String()).
Str("kind", job.Kind).
Int64("attempt", job.Attempts).
Msg("running job")
start := time.Now()
runErr := handler(ctx, q, job)
elapsed := time.Since(start)
if runErr == nil {
l.Info().
Str("job_id", job.ID.String()).
Str("kind", job.Kind).
Dur("duration_ms", elapsed).
Msg("job succeeded")
_, err = q.MarkJobSucceeded(ctx, sqlc.MarkJobSucceededParams{
ID: job.ID,
})
return err
}
l.Warn().
Str("job_id", job.ID.String()).
Str("kind", job.Kind).
Int64("attempt", job.Attempts).
Dur("duration_ms", elapsed).
Err(runErr).
Msg("job failed")
if job.Attempts >= job.MaxAttempts {
msg := runErr.Error()
_, _ = q.MarkJobFailed(ctx, sqlc.MarkJobFailedParams{
LastError: &msg,
ID: job.ID,
})
return nil
}
backoff := opts.backoff(int(job.Attempts))
nextRun := opts.now().Add(backoff)
msg := runErr.Error()
_, err = q.RequeueJob(ctx, sqlc.RequeueJobParams{
ID: job.ID,
RunAt: nextRun,
LastError: &msg,
})
return err
}
// RunLoop polls for work until ctx is canceled.
func RunLoop(ctx context.Context, q *sqlc.Queries, opts *Options) error {
if opts == nil {
opts = &Options{}
}
l := opts.logger()
sleep := opts.sleep()
for {
if ctx.Err() != nil {
return ctx.Err()
}
if err := RunOnce(ctx, q, opts); err != nil {
l.Error().Err(err).Msg("worker run once failed")
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(sleep):
}
continue
}
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(sleep):
}
}
}