llmgw/internal/web/download_test.go

package web

import (
	"archive/zip"
	"bytes"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func zipOf(t *testing.T, reply string) (*zip.Reader, *httptest.ResponseRecorder) {
	t.Helper()
	body := `{"content":` + jsonString(reply) + `,"name":"code"}`
	r := httptest.NewRequest(http.MethodPost, "http://x/chat/download", strings.NewReader(body))
	w := httptest.NewRecorder()
	CodeArchive(1<<20)(w, r)
	if w.Code != http.StatusOK {
		return nil, w
	}
	zr, err := zip.NewReader(bytes.NewReader(w.Body.Bytes()), int64(w.Body.Len()))
	if err != nil {
		t.Fatalf("the response is not a readable zip: %v", err)
	}
	return zr, w
}

func jsonString(s string) string {
	var b strings.Builder
	b.WriteByte('"')
	for _, r := range s {
		switch r {
		case '"':
			b.WriteString(`\"`)
		case '\\':
			b.WriteString(`\\`)
		case '\n':
			b.WriteString(`\n`)
		case '\t':
			b.WriteString(`\t`)
		default:
			b.WriteRune(r)
		}
	}
	b.WriteByte('"')
	return b.String()
}

func TestTheArchiveKeepsDirectories(t *testing.T) {
	// The point of the feature: a reply is often a small tree, not one file.
	zr, _ := zipOf(t, "```go:cmd/api/main.go\npackage main\n```\n"+
		"```go:internal/store/store.go\npackage store\n```\n"+
		"```yaml:deploy/k8s/deployment.yaml\nkind: Deployment\n```\n")
	got := map[string]string{}
	for _, f := range zr.File {
		rc, err := f.Open()
		if err != nil {
			t.Fatal(err)
		}
		b, _ := io.ReadAll(rc)
		rc.Close()
		got[f.Name] = string(b)
	}
	for _, want := range []string{"cmd/api/main.go", "internal/store/store.go", "deploy/k8s/deployment.yaml"} {
		if _, ok := got[want]; !ok {
			t.Errorf("missing %s; archive has %v", want, keys(got))
		}
	}
	if !strings.Contains(got["cmd/api/main.go"], "package main") {
		t.Errorf("contents lost: %q", got["cmd/api/main.go"])
	}
}

func keys(m map[string]string) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	return out
}

func TestNoEntryEscapesTheArchiveRoot(t *testing.T) {
	// Every name in this archive was written by a language model, and someone
	// is going to extract it with a tool that trusts the archive.
	zr, _ := zipOf(t, "```go:../../.ssh/authorized_keys\nssh-rsa AAAA\n```\n"+
		"```go:/etc/cron.d/x\n* * * * * root sh\n```\n"+
		"```go:a/../../../b.go\npackage b\n```\n")
	if zr == nil {
		t.Fatal("nothing was produced")
	}
	for _, f := range zr.File {
		if strings.HasPrefix(f.Name, "/") || strings.HasPrefix(f.Name, "../") ||
			strings.Contains(f.Name, "/../") || strings.Contains(f.Name, `\`) {
			t.Errorf("entry escapes the root: %q", f.Name)
		}
	}
}

func TestAReplyWithNoCodeIsRefusedRatherThanEmpty(t *testing.T) {
	// An empty zip downloads successfully and tells the reader nothing.
	_, w := zipOf(t, "There is no code here, only prose.")
	if w.Code != http.StatusBadRequest {
		t.Errorf("status %d; a reply with no code should be refused", w.Code)
	}
}

func TestTheDownloadIsSentAsAnAttachment(t *testing.T) {
	_, w := zipOf(t, "```go:main.go\npackage main\n```\n")
	if ct := w.Header().Get("Content-Type"); ct != "application/zip" {
		t.Errorf("Content-Type %q", ct)
	}
	if cd := w.Header().Get("Content-Disposition"); !strings.HasPrefix(cd, "attachment;") {
		t.Errorf("Content-Disposition %q; it must not render in the tab", cd)
	}
	if w.Header().Get("X-Content-Type-Options") != "nosniff" {
		t.Error("the archive may be sniffed as another type")
	}
}

func TestTheArchiveNameIsNotTakenOnTrust(t *testing.T) {
	for in, want := range map[string]string{
		"code":             "code.zip",
		"":                 "code.zip",
		"../../etc/passwd": "etcpasswd.zip",
		// Hyphens are kept on purpose, so "rm -rf" survives as "rm-rf". The
		// point of this case is the quotes and slashes, which do not.
		`a"; rm -rf /; x="`:      "arm-rfx.zip",
		strings.Repeat("x", 100): strings.Repeat("x", 40) + ".zip",
	} {
		if got := archiveName(in); got != want {
			t.Errorf("archiveName(%q) = %q, want %q", in, got, want)
		}
	}
}