@pi-in-go/pigpen-ahp 0.1.0 → 0.1.1

package/extensions/ahp/internal/pi/chatdriver.go

@@ -95,7 +95,10 @@ func (d *ChatDriver) dispatchAll(actions []ahptypes.StateAction) {
 	}
 }
 
-// track runs fn as an in-flight operation that Quiesce waits for.
+// track runs fn as an in-flight operation that Quiesce waits for. These goroutines carry client
+// input into the backend (prompt and steering text and images, decoded there), outside the
+// host's request recover, so a panic is logged here instead of ending the host process; the
+// prompt and steering paths turn it into their usual error first (callBackend).
 func (d *ChatDriver) track(fn func()) {
 	d.mu.Lock()
 	d.inflight++
@@ -107,10 +110,26 @@ func (d *ChatDriver) track(fn func()) {
 			d.idle.Broadcast()
 			d.mu.Unlock()
 		}()
+		defer func() {
+			if p := recover(); p != nil {
+				d.logf("chat operation failed: %v", p)
+			}
+		}()
 		fn()
 	}()
 }
 
+// callBackend runs one backend call and returns a panic in it as an error, so a client message
+// that breaks the backend (a hostile image in its decoders) fails like a rejected prompt.
+func callBackend(fn func() error) (err error) {
+	defer func() {
+		if p := recover(); p != nil {
+			err = fmt.Errorf("internal error: %v", p)
+		}
+	}()
+	return fn()
+}
+
 func (d *ChatDriver) drainInFlight() {
 	d.mu.Lock()
 	for d.inflight > 0 {
@@ -222,10 +241,13 @@ func (d *ChatDriver) steer(ctx context.Context, message ahptypes.Message) {
 	if ctx.Err() != nil {
 		return
 	}
-	input, err := messageInput(message)
-	if err == nil {
-		err = d.opts.Backend.Steer(ctx, input.Text, input.Images)
-	}
+	err := callBackend(func() error {
+		input, err := messageInput(message)
+		if err != nil {
+			return err
+		}
+		return d.opts.Backend.Steer(ctx, input.Text, input.Images)
+	})
 	if err != nil {
 		d.logf("steer failed: %v", err)
 		d.clearPendingSteering()
@@ -277,10 +299,13 @@ func (d *ChatDriver) runPrompt(m *mapper.TurnMapper, message ahptypes.Message, l
 		if !d.stillCurrent(m, ctx) {
 			return
 		}
-		input, err := messageInput(message)
-		if err == nil {
-			err = d.opts.Backend.Prompt(ctx, input.Text, input.Images)
-		}
+		err := callBackend(func() error {
+			input, err := messageInput(message)
+			if err != nil {
+				return err
+			}
+			return d.opts.Backend.Prompt(ctx, input.Text, input.Images)
+		})
 		if err != nil {
 			d.mu.Lock()
 			current := d.mapper == m

package/extensions/ahp/internal/pi/imagebomb_test.go

@@ -0,0 +1,89 @@
+package pi
+
+import (
+	"bytes"
+	"compress/zlib"
+	"encoding/base64"
+	"encoding/binary"
+	"hash/crc32"
+	"runtime"
+	"strings"
+	"testing"
+
+	"github.com/MichaelKinsy/pigpen/ahp/internal/mapper"
+)
+
+// bombPNG streams a w×h 8-bit grayscale PNG of zeros through zlib without ever holding the
+// pixels, so the test itself stays small: about 400 KiB for 20000×20000.
+func bombPNG(t *testing.T, w, h int) []byte {
+	t.Helper()
+	var out bytes.Buffer
+	out.WriteString("\x89PNG\r\n\x1a\n")
+	chunk := func(kind string, body []byte) {
+		var n [4]byte
+		binary.BigEndian.PutUint32(n[:], uint32(len(body)))
+		out.Write(n[:])
+		out.WriteString(kind)
+		out.Write(body)
+		crc := crc32.NewIEEE()
+		crc.Write([]byte(kind))
+		crc.Write(body)
+		binary.BigEndian.PutUint32(n[:], crc.Sum32())
+		out.Write(n[:])
+	}
+	ihdr := make([]byte, 13)
+	binary.BigEndian.PutUint32(ihdr[0:], uint32(w))
+	binary.BigEndian.PutUint32(ihdr[4:], uint32(h))
+	ihdr[8] = 8 // bit depth; colour type 0 is grayscale
+	chunk("IHDR", ihdr)
+	var idat bytes.Buffer
+	// BestSpeed: about the same size for all-zero rows, and several times faster (this runs under -race too).
+	zw, err := zlib.NewWriterLevel(&idat, zlib.BestSpeed)
+	if err != nil {
+		t.Fatal(err)
+	}
+	row := make([]byte, 1+w) // filter byte 0 plus zero pixels
+	for y := 0; y < h; y++ {
+		if _, err := zw.Write(row); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := zw.Close(); err != nil {
+		t.Fatal(err)
+	}
+	chunk("IDAT", idat.Bytes())
+	chunk("IEND", nil)
+	return out.Bytes()
+}
+
+// A client-supplied 20000×20000 PNG (400 MP, about 400 KiB on the wire) must be refused from its
+// header, before the decoder allocates the pixel buffer.
+func TestPrepareImagesRefusesDecodeBomb(t *testing.T) {
+	raw := bombPNG(t, 20000, 20000)
+	if len(raw) > 2<<20 {
+		t.Fatalf("fixture is %d KiB, expected a small bomb", len(raw)>>10)
+	}
+	data := base64.StdEncoding.EncodeToString(raw)
+	for _, autoResize := range []bool{true, false} {
+		runtime.GC()
+		var before, after runtime.MemStats
+		runtime.ReadMemStats(&before)
+		_, err := PrepareImagesForPi([]mapper.Image{{Type: "image", Data: data, MimeType: "image/png"}}, autoResize)
+		runtime.ReadMemStats(&after)
+		if autoResize && (err == nil || !strings.Contains(err.Error(), "too large")) {
+			t.Fatalf("autoResize=%v: want a too-large refusal, got %v", autoResize, err)
+		}
+		if grown := after.TotalAlloc - before.TotalAlloc; grown > 32<<20 {
+			t.Fatalf("autoResize=%v: refusing allocated %d MiB (want under 32)", autoResize, grown>>20)
+		}
+	}
+}
+
+// Images at or under the cap still work.
+func TestPrepareImagesAcceptsLargeButBoundedPNG(t *testing.T) {
+	data := base64.StdEncoding.EncodeToString(bombPNG(t, 3000, 3000))
+	out, err := PrepareImagesForPi([]mapper.Image{{Type: "image", Data: data, MimeType: "image/png"}}, true)
+	if err != nil || len(out) != 1 {
+		t.Fatalf("%v %v", out, err)
+	}
+}

package/extensions/ahp/internal/pi/imageinput.go

@@ -24,6 +24,11 @@ import (
 const (
 	maxImageDimension = 2000
 	maxImageBytes     = 4_500_000
+	// maxDecodePixels bounds what a client image may declare before it is decoded: the decoders
+	// allocate the whole pixel buffer from the header, so 420 KiB of PNG can ask for 20000×20000
+	// pixels (hundreds of MiB, seconds of CPU). Same value and name as websearch's imageproc.go
+	// (a separate module; websearch's TestDecodePixelCapMatchesAHP keeps them equal).
+	maxDecodePixels = 1 << 26
 )
 
 var supportedInlineTypes = map[string]bool{"image/png": true, "image/jpeg": true, "image/gif": true, "image/webp": true}
@@ -68,7 +73,25 @@ func prepareImage(img mapper.Image, autoResize bool, index int) (mapper.Image, e
 	return mapper.Image{Type: "image", Data: resized, MimeType: resizedType}, nil
 }
 
+// checkDecodeSize reads only the header and refuses images declaring more than maxDecodePixels.
+func checkDecodeSize(raw []byte) error {
+	cfg, _, err := image.DecodeConfig(bytes.NewReader(raw))
+	if err != nil {
+		return err
+	}
+	if cfg.Width <= 0 || cfg.Height <= 0 || int64(cfg.Width)*int64(cfg.Height) > maxDecodePixels {
+		return fmt.Errorf("image too large to process (%d×%d pixels; the limit is %d megapixels)", cfg.Width, cfg.Height, maxDecodePixels/(1<<20))
+	}
+	return nil
+}
+
 func decodeAny(raw []byte, mimeType string) (image.Image, error) {
+	switch mimeType {
+	case "image/png", "image/jpeg", "image/gif":
+		if err := checkDecodeSize(raw); err != nil {
+			return nil, err
+		}
+	}
 	switch mimeType {
 	case "image/png":
 		return png.Decode(bytes.NewReader(raw))

package/extensions/ahp/internal/pi/panic_test.go

@@ -0,0 +1,44 @@
+package pi_test
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/microsoft/agent-host-protocol/clients/go/ahptypes"
+
+	"github.com/MichaelKinsy/pigpen/ahp/internal/testkit"
+)
+
+// The chat driver runs a client's prompt and steering on goroutines of its own (ChatDriver.track),
+// and that is where client images are decoded (live.Session → PrepareImagesForPi). A panic there
+// was outside every recover in the host, so one hostile message ended the AHP host process. It
+// now ends that turn (or drops that steering message) with an error, as a rejected prompt does.
+
+func TestChatDriverSurvivesPanickingPrompt(t *testing.T) {
+	backend := funcBackend{prompt: func(context.Context, string) error { panic("decoder exploded") }}
+	f := startFixture(t, backend)
+	f.turnStarted("t1", userMessage("hi"))
+	testkit.Eventually(t, "the panicking prompt to close its turn", turnCount(f, 1))
+	turn := f.chat().Turns[0]
+	if turn.State != ahptypes.TurnStateError {
+		t.Fatalf("state = %s", turn.State)
+	}
+	last, ok := turn.ResponseParts[len(turn.ResponseParts)-1].Value.(*ahptypes.ErrorResponsePart)
+	if !ok || !strings.Contains(last.Error.Message, "decoder exploded") {
+		t.Fatalf("error part = %+v", turn.ResponseParts)
+	}
+	// The host still serves the next turn.
+	f.turnStarted("t2", userMessage("again"))
+	testkit.Eventually(t, "the next turn to close too", turnCount(f, 2))
+}
+
+func TestChatDriverSurvivesPanickingSteer(t *testing.T) {
+	backend := funcBackend{steer: func(context.Context, string) error { panic("decoder exploded") }}
+	f := startFixture(t, backend)
+	f.turnStarted("t1", userMessage("long running"))
+	testkit.Eventually(t, "the turn to be active", func() bool { return f.chat().ActiveTurn != nil })
+	f.pending("steering", "steer-1", userMessage("focus"))
+	testkit.Eventually(t, "the steering message to enter protocol state", func() bool { return f.chat().SteeringMessage != nil })
+	testkit.Eventually(t, "the failed steering message to leave protocol state", func() bool { return f.chat().SteeringMessage == nil })
+}

package/extensions/ahp/internal/pi/summary_test.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"reflect"
 	"testing"
+	"time"
 
 	"github.com/microsoft/agent-host-protocol/clients/go/ahp"
 	"github.com/microsoft/agent-host-protocol/clients/go/ahptypes"
@@ -357,6 +358,19 @@ func projectsOutcome(t *testing.T, outcome string) {
 	}
 }
 
+// waitUntil polls condition until it holds. The bound only turns a lost projection into a failure
+// instead of a hang; a passing run does not depend on it.
+func waitUntil(t *testing.T, what string, condition func() bool) {
+	t.Helper()
+	deadline := time.Now().Add(30 * time.Second)
+	for !condition() {
+		if time.Now().After(deadline) {
+			t.Fatalf("timed out waiting for %s", what)
+		}
+		time.Sleep(time.Millisecond)
+	}
+}
+
 func replaysProjection(t *testing.T, running bool) {
 	f := newSummaryFixture(t, nil)
 	baseline := f.h.host.ServerSeq()
@@ -372,6 +386,19 @@ func replaysProjection(t *testing.T, running bool) {
 		h.DispatchServerAction(f.chat, ahptypes.StateAction{Value: &ahptypes.ChatActivityChangedAction{Type: ahptypes.ActionTypeChatActivityChanged}})
 		h.DispatchServerAction(f.chat, ahptypes.StateAction{Value: &ahptypes.ChatTurnCompleteAction{Type: ahptypes.ActionTypeChatTurnComplete, TurnId: "turn", Duration: 1000}})
 	}
+	// DispatchServerAction returns once its own listener events are queued: when another goroutine
+	// is already draining the host's listener queue (an earlier action's projection), the call
+	// returns before the chat-to-session projection of these actions has run. A reconnect taken
+	// at that instant replays a session channel that lacks the projection, so the mirror keeps the
+	// chat's creation time. Wait for the projection the way a client would see it.
+	wantModified := endAt
+	if running {
+		wantModified = startAt
+	}
+	waitUntil(t, "the session summary to project the chat's modification time", func() bool {
+		session := f.sessionState()
+		return session != nil && len(session.Chats) == 1 && session.Chats[0].ModifiedAt == wantModified
+	})
 	resumed := testkit.Connect(t, h)
 	t.Cleanup(resumed.Close)
 	var replay struct {
@@ -402,10 +429,6 @@ func replaysProjection(t *testing.T, running bool) {
 	if mirror.Status != ahptypes.SessionStatusIdle {
 		t.Fatalf("mirror status = %d", mirror.Status)
 	}
-	wantModified := endAt
-	if running {
-		wantModified = startAt
-	}
 	if mirror.Chats[0].ModifiedAt != wantModified {
 		t.Fatalf("chat modifiedAt = %s", mirror.Chats[0].ModifiedAt)
 	}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pi-in-go/pigpen-ahp",
-  "version": "0.1.0",
+  "version": "0.1.1",
   "description": "Go port of pi-ahp: serves the running PiG session to remote Microsoft Agent Host Protocol clients over WebSocket, with reconnect and auth. No listener unless you ask; no remote filesystem or terminal unless you opt in",
   "keywords": [
     "pig-package",

package/proof/PORT.md

@@ -96,7 +96,9 @@ session. The equivalence method is therefore three layers, not the scenario harn
    service refuses to start there. A signalled child reports `128+signal`.
 9. **Images.** Pi's `convertToPng`/`resizeImage` are not in the SDK: PNG, JPEG, GIF decode with the
    standard library, BMP with a small decoder, resizing is a box filter, WebP passes through unvalidated
-   beyond its header.
+   beyond its header. PNG, JPEG and GIF declaring more than 64 megapixels are refused from their header,
+   before decoding (Go's decoders allocate the whole image from the header). A
+   panic while a prompt or steering message is handed to the agent fails that turn or message.
 10. **MIME table** is the subset of the `mime` package that matters for source trees.
 11. **Async side effects.** Registry and driver effects run on goroutines where upstream runs them in
     the same tick; a few twins wait for the backend call before the next step. `host.drain` defers