diff --git a/client.go b/client.go index c04f1afd..465c4beb 100644 --- a/client.go +++ b/client.go @@ -397,15 +397,25 @@ func (c *Client) ReadDirContext(ctx context.Context, p string) ([]os.FileInfo, e } switch typ { case sshFxpName: - sid, data := unmarshalUint32(data) - if sid != id { - return nil, &unexpectedIDErr{id, sid} + data, err = unmarshalSID(id, data) + if err != nil { + return nil, err + } + var count int + count, data, err = unmarshalCount(data) + if err != nil { + return nil, err } - count, data := unmarshalUint32(data) - for i := uint32(0); i < count; i++ { + for i := 0; i < count; i++ { var filename string - filename, data = unmarshalString(data) - _, data = unmarshalString(data) // discard longname + filename, data, err = unmarshalStringSafe(data) + if err != nil { + return nil, err + } + _, data, err = unmarshalStringSafe(data) // discard longname + if err != nil { + return nil, err + } var attr *FileStat attr, data, err = unmarshalAttrs(data) if err != nil { @@ -441,11 +451,14 @@ func (c *Client) opendir(ctx context.Context, path string) (string, error) { } switch typ { case sshFxpHandle: - sid, data := unmarshalUint32(data) - if sid != id { - return "", &unexpectedIDErr{id, sid} + data, err := unmarshalSID(id, data) + if err != nil { + return "", err + } + handle, _, err := unmarshalStringSafe(data) + if err != nil { + return "", err } - handle, _ := unmarshalString(data) return handle, nil case sshFxpStatus: return "", normaliseError(unmarshalStatus(id, data)) @@ -477,9 +490,9 @@ func (c *Client) Lstat(p string) (os.FileInfo, error) { } switch typ { case sshFxpAttrs: - sid, data := unmarshalUint32(data) - if sid != id { - return nil, &unexpectedIDErr{id, sid} + data, err := unmarshalSID(id, data) + if err != nil { + return nil, err } attr, _, err := unmarshalAttrs(data) if err != nil { @@ -506,15 +519,21 @@ func (c *Client) ReadLink(p string) (string, error) { } switch typ { case sshFxpName: - sid, data := unmarshalUint32(data) - if sid != id { - return "", &unexpectedIDErr{id, sid} + data, err := unmarshalSID(id, data) + if err != nil { + return "", err + } + count, data, err := unmarshalCount(data) + if err != nil { + return "", err } - count, data := unmarshalUint32(data) if count != 1 { - return "", unexpectedCount(1, count) + return "", unexpectedCount(1, uint32(count)) + } + filename, _, err := unmarshalStringSafe(data) // ignore dummy attributes + if err != nil { + return "", err } - filename, _ := unmarshalString(data) // ignore dummy attributes return filename, nil case sshFxpStatus: return "", normaliseError(unmarshalStatus(id, data)) @@ -677,11 +696,14 @@ func (c *Client) open(path string, pflags uint32) (*File, error) { } switch typ { case sshFxpHandle: - sid, data := unmarshalUint32(data) - if sid != id { - return nil, &unexpectedIDErr{id, sid} + data, err := unmarshalSID(id, data) + if err != nil { + return nil, err + } + handle, _, err := unmarshalStringSafe(data) + if err != nil { + return nil, err } - handle, _ := unmarshalString(data) return &File{c: c, path: path, handle: handle}, nil case sshFxpStatus: return nil, normaliseError(unmarshalStatus(id, data)) @@ -721,9 +743,9 @@ func (c *Client) stat(path string) (*FileStat, error) { } switch typ { case sshFxpAttrs: - sid, data := unmarshalUint32(data) - if sid != id { - return nil, &unexpectedIDErr{id, sid} + data, err := unmarshalSID(id, data) + if err != nil { + return nil, err } attr, _, err := unmarshalAttrs(data) return attr, err @@ -745,9 +767,9 @@ func (c *Client) fstat(handle string) (*FileStat, error) { } switch typ { case sshFxpAttrs: - sid, data := unmarshalUint32(data) - if sid != id { - return nil, &unexpectedIDErr{id, sid} + data, err := unmarshalSID(id, data) + if err != nil { + return nil, err } attr, _, err := unmarshalAttrs(data) return attr, err @@ -942,15 +964,21 @@ func (c *Client) RealPath(path string) (string, error) { } switch typ { case sshFxpName: - sid, data := unmarshalUint32(data) - if sid != id { - return "", &unexpectedIDErr{id, sid} + data, err := unmarshalSID(id, data) + if err != nil { + return "", err + } + count, data, err := unmarshalCount(data) + if err != nil { + return "", err } - count, data := unmarshalUint32(data) if count != 1 { - return "", unexpectedCount(1, count) + return "", unexpectedCount(1, uint32(count)) + } + filename, _, err := unmarshalStringSafe(data) // ignore attributes + if err != nil { + return "", err } - filename, _ := unmarshalString(data) // ignore attributes return filename, nil case sshFxpStatus: return "", normaliseError(unmarshalStatus(id, data)) @@ -1147,13 +1175,11 @@ func (f *File) readChunkAt(ch chan result, b []byte, off int64) (n int, err erro return n, normaliseError(unmarshalStatus(id, data)) case sshFxpData: - sid, data := unmarshalUint32(data) - if id != sid { - return n, &unexpectedIDErr{id, sid} + data, err = unmarshalDataReply(id, data) + if err != nil { + return n, err } - - l, data := unmarshalUint32(data) - n += copy(b[n:], data[:l]) + n += copy(b[n:], data) default: return n, unimplementedPacketErr(typ) @@ -1293,13 +1319,10 @@ func (f *File) readAt(b []byte, off int64) (int, error) { err = normaliseError(unmarshalStatus(packet.id, s.data)) case sshFxpData: - sid, data := unmarshalUint32(s.data) - if packet.id != sid { - err = &unexpectedIDErr{packet.id, sid} - - } else { - l, data := unmarshalUint32(data) - n = copy(packet.b, data[:l]) + var data []byte + data, err = unmarshalDataReply(packet.id, s.data) + if err == nil { + n = copy(packet.b, data) // For normal disk files, it is guaranteed that this will read // the specified number of bytes, or up to end of file. @@ -1522,14 +1545,15 @@ func (f *File) WriteTo(w io.Writer) (written int64, err error) { err = normaliseError(unmarshalStatus(readWork.id, s.data)) case sshFxpData: - sid, data := unmarshalUint32(s.data) - if readWork.id != sid { - err = &unexpectedIDErr{readWork.id, sid} - - } else { - l, data := unmarshalUint32(data) - b = pool.Get()[:l] - n = copy(b, data[:l]) + var data []byte + data, err = unmarshalDataReply(readWork.id, s.data) + if err == nil { + // A malicious server may return more data than we + // requested (up to the packet limit, larger than the + // pool buffer); copy bounds it to the buffer size + // rather than slicing past the page and panicking. + b = pool.Get() + n = copy(b, data) b = b[:n] } @@ -1645,8 +1669,11 @@ func (f *File) writeChunkAt(ch chan result, b []byte, off int64) (int, error) { switch typ { case sshFxpStatus: - id, _ := unmarshalUint32(data) - err := normaliseError(unmarshalStatus(id, data)) + id, _, err := unmarshalUint32Safe(data) + if err != nil { + return 0, err + } + err = normaliseError(unmarshalStatus(id, data)) if err != nil { return 0, err } diff --git a/client_test.go b/client_test.go index dda8af2b..0b26762a 100644 --- a/client_test.go +++ b/client_test.go @@ -181,6 +181,272 @@ func TestClientShortPacket(t *testing.T) { } } +// frameReply returns a server reply (length prefix + type + body, where body +// begins with the echoed request id) as recvPacket expects it on the wire. +func frameReply(typ fxp, body []byte) []byte { + b := append([]byte{byte(typ)}, body...) + return append(marshalUint32(nil, uint32(len(b))), b...) +} + +// readRequest reads one length-prefixed request packet from r and returns its +// type and request id. The SSH_FXP_INIT handshake packet carries a version in +// place of an id; callers special-case it on the type. +func readRequest(r io.Reader) (typ fxp, id uint32, ok bool) { + hdr := make([]byte, 4) + if _, err := io.ReadFull(r, hdr); err != nil { + return 0, 0, false + } + n, _ := unmarshalUint32(hdr) + body := make([]byte, n) + if _, err := io.ReadFull(r, body); err != nil { + return 0, 0, false + } + typ = fxp(body[0]) + if len(body) >= 5 { + id, _ = unmarshalUint32(body[1:]) + } + return typ, id, true +} + +// startScriptedClient wires a *Client to a scripted server over synchronized +// io.Pipes. After answering the SSH_FXP_INIT/VERSION handshake, the server +// replies to each request by calling reply(typ, id); reply returns the response +// packet type and body (which must begin with the echoed id). Returning +// ok==false stops the server. +func startScriptedClient(t *testing.T, reply func(typ fxp, id uint32) (fxp, []byte, bool), opts ...ClientOption) *Client { + t.Helper() + + cr, sw := io.Pipe() // client reads what the server writes + sr, cw := io.Pipe() // server reads what the client writes + + go func() { + defer sw.Close() + for { + typ, id, ok := readRequest(sr) + if !ok { + return + } + if typ == sshFxpInit { + sendPacket(sw, &sshFxVersionPacket{Version: sftpProtocolVersion}) + continue + } + rtyp, body, ok := reply(typ, id) + if !ok { + return + } + if _, err := sw.Write(frameReply(rtyp, body)); err != nil { + return + } + } + }() + + c, err := NewClientPipe(cr, cw, opts...) + if err != nil { + t.Fatal(err) + } + return c +} + +func TestClientReadDataReply(t *testing.T) { + const payload = "ABCDEFGH" + + cases := []struct { + name string + data func(b []byte) []byte + wantErr error // nil means the read must succeed and return payload + }{ + {"valid", func(b []byte) []byte { return marshalString(b, payload) }, nil}, + {"length exceeds bytes present", func(b []byte) []byte { return marshalUint32(b, 0x7fffffff) }, errShortPacket}, + {"length prefix truncated", func(b []byte) []byte { return append(b, 0x00, 0x00) }, errShortPacket}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c := startScriptedClient(t, func(_ fxp, id uint32) (fxp, []byte, bool) { + return sshFxpData, tc.data(marshalUint32(nil, id)), true + }) + defer c.Close() + + f := &File{c: c, path: "somefile", handle: "somehandle"} + b := make([]byte, len(payload)) + n, err := f.Read(b) + + if tc.wantErr != nil { + if !errors.Is(err, tc.wantErr) { + t.Fatalf("error = %v, want %v", err, tc.wantErr) + } + if n != 0 { + t.Errorf("n = %d, want 0", n) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if string(b[:n]) != payload { + t.Errorf("read %q, want %q", b[:n], payload) + } + }) + } +} + +func TestClientReadAtConcurrentDataReply(t *testing.T) { + c := startScriptedClient(t, func(_ fxp, id uint32) (fxp, []byte, bool) { + // Oversize declared length, no data. + body := marshalUint32(nil, id) + body = marshalUint32(body, 0x7fffffff) + return sshFxpData, body, true + }, MaxPacket(1024)) + defer c.Close() + + f := &File{c: c, path: "somefile", handle: "somehandle"} + // A buffer larger than maxPacket forces readAt to split the read into + // concurrent maxPacket-sized requests. + n, err := f.ReadAt(make([]byte, 4096), 0) + if !errors.Is(err, errShortPacket) { + t.Fatalf("error = %v, want %v", err, errShortPacket) + } + if n != 0 { + t.Errorf("n = %d, want 0", n) + } +} + +// dirEntry appends one well-formed SSH_FXP_NAME entry (filename, longname, and +// empty attributes) to b. +func dirEntry(b []byte, name string) []byte { + b = marshalString(b, name) // filename + b = marshalString(b, name) // longname + return marshalUint32(b, 0) // attrs: flags == 0, no fields follow +} + +// statusBody appends an SSH_FXP_STATUS body (code plus empty message/language) +// after an already-written request id. +func statusBody(b []byte, code uint32) []byte { + b = marshalUint32(b, code) + b = marshalString(b, "") // message + return marshalString(b, "") // language +} + +func TestClientReadDirMalformedName(t *testing.T) { + // body builds the SSH_FXP_NAME reply body that follows the echoed id. (The + // 32-bit overflow of the count field is covered by TestUnmarshalCount.) + cases := []struct { + name string + body func(b []byte) []byte + wantErr error + }{ + { + name: "count prefix truncated", + body: func(b []byte) []byte { return append(b, 0x00, 0x00) }, + wantErr: errShortPacket, + }, + { + // A huge count with no entries must fail when the loop runs out of + // bytes, not pre-allocate or panic. + name: "count far exceeds entries present", + body: func(b []byte) []byte { return marshalUint32(b, 0x7fffffff) }, + wantErr: errShortPacket, + }, + { + name: "filename length exceeds bytes present", + body: func(b []byte) []byte { + b = marshalUint32(b, 1) // count: one entry + return marshalUint32(b, 0x7fffffff) // filename length, no bytes follow + }, + wantErr: errShortPacket, + }, + { + name: "longname length exceeds bytes present", + body: func(b []byte) []byte { + b = marshalUint32(b, 1) // count: one entry + b = marshalString(b, "file") // valid filename + return marshalUint32(b, 0x7fffffff) // longname length, no bytes follow + }, + wantErr: errShortPacket, + }, + { + name: "attrs truncated", + body: func(b []byte) []byte { + b = marshalUint32(b, 1) // count: one entry + b = marshalString(b, "file") // filename + b = marshalString(b, "file") // longname + // Attr flags claim a size field, but the 8 bytes never follow. + return marshalUint32(b, sshFileXferAttrSize) + }, + wantErr: errShortPacket, + }, + { + name: "second entry truncated mid-loop", + body: func(b []byte) []byte { + b = marshalUint32(b, 2) // declare two entries + b = dirEntry(b, "file1") // first entry is well-formed + return marshalUint32(b, 0x7fffffff) // second filename length, no bytes + }, + wantErr: errShortPacket, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c := startScriptedClient(t, func(typ fxp, id uint32) (fxp, []byte, bool) { + switch typ { + case sshFxpOpendir: + return sshFxpHandle, marshalString(marshalUint32(nil, id), "somehandle"), true + case sshFxpReaddir: + return sshFxpName, tc.body(marshalUint32(nil, id)), true + default: // the deferred close, etc. + return 0, nil, false + } + }) + defer c.Close() + + if _, err := c.ReadDir("/"); !errors.Is(err, tc.wantErr) { + t.Fatalf("ReadDir error = %v, want %v", err, tc.wantErr) + } + }) + } +} + +// TestClientReadDirSuccess is the positive counterpart: a well-formed listing +// (two entries, then an EOF status) must be parsed into the expected entries +// with no error, proving the hardened loop still accepts valid replies. +func TestClientReadDirSuccess(t *testing.T) { + var readdirs int + c := startScriptedClient(t, func(typ fxp, id uint32) (fxp, []byte, bool) { + switch typ { + case sshFxpOpendir: + return sshFxpHandle, marshalString(marshalUint32(nil, id), "somehandle"), true + case sshFxpReaddir: + readdirs++ + if readdirs > 1 { + // End of listing on the second readdir. + return sshFxpStatus, statusBody(marshalUint32(nil, id), sshFxEOF), true + } + body := marshalUint32(nil, id) + body = marshalUint32(body, 2) // two entries + body = dirEntry(body, "file1") + body = dirEntry(body, "file2") + return sshFxpName, body, true + case sshFxpClose: + return sshFxpStatus, statusBody(marshalUint32(nil, id), sshFxOk), true + default: + return 0, nil, false + } + }) + defer c.Close() + + entries, err := c.ReadDir("/") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(entries) != 2 { + t.Fatalf("got %d entries, want 2", len(entries)) + } + if entries[0].Name() != "file1" || entries[1].Name() != "file2" { + t.Errorf("entries = [%q %q], want [file1 file2]", entries[0].Name(), entries[1].Name()) + } +} + // Issue #418: panic in clientConn.recv when the sid is incomplete. func TestClientNoSid(t *testing.T) { stream := new(bytes.Buffer) diff --git a/packet.go b/packet.go index 74712825..87751ac0 100644 --- a/packet.go +++ b/packet.go @@ -164,11 +164,11 @@ func unmarshalString(b []byte) (string, []byte) { } func unmarshalStringSafe(b []byte) (string, []byte, error) { - n, b, err := unmarshalUint32Safe(b) + n, b, err := unmarshalCount(b) if err != nil { return "", nil, err } - if int64(n) > int64(len(b)) { + if n > len(b) { return "", nil, errShortPacket } return string(b[:n]), b[n:], nil @@ -253,12 +253,44 @@ func unmarshalFileStat(flags uint32, b []byte) (*FileStat, []byte, error) { return &fs, b, nil } -func unmarshalStatus(id uint32, data []byte) error { - sid, data := unmarshalUint32(data) +// unmarshalSID consumes the request-id echoed back in a reply and verifies it +// matches the id we sent, returning the remaining bytes for further decoding. +func unmarshalSID(id uint32, b []byte) ([]byte, error) { + sid, b, err := unmarshalUint32Safe(b) + if err != nil { + return nil, err + } if sid != id { - return &unexpectedIDErr{id, sid} + return nil, &unexpectedIDErr{id, sid} + } + return b, nil +} + +// unmarshalCount consumes a uint32 count and returns it as an int. A value that +// does not fit in a non-negative int is rejected with errLongPacket: this is +// only possible when strconv.IntSize == 32 and the count exceeds math.MaxInt32, +// in which case it is longer than any possible slice and would otherwise drive a +// negative length into make/slice and panic. +func unmarshalCount(b []byte) (int, []byte, error) { + v, b, err := unmarshalUint32Safe(b) + if err != nil { + return 0, nil, err + } + if int(v) < 0 { + return 0, nil, errLongPacket + } + return int(v), b, nil +} + +func unmarshalStatus(id uint32, data []byte) error { + data, err := unmarshalSID(id, data) + if err != nil { + return err + } + code, data, err := unmarshalUint32Safe(data) + if err != nil { + return err } - code, data := unmarshalUint32(data) msg, data, _ := unmarshalStringSafe(data) lang, _, _ := unmarshalStringSafe(data) return &StatusError{ @@ -268,6 +300,24 @@ func unmarshalStatus(id uint32, data []byte) error { } } +// unmarshalDataReply parses an SSH_FXP_DATA reply, validating the id and the +// declared data length against the bytes actually present. It returns the data +// payload as a sub-slice of data. +func unmarshalDataReply(id uint32, data []byte) ([]byte, error) { + data, err := unmarshalSID(id, data) + if err != nil { + return nil, err + } + l, data, err := unmarshalCount(data) + if err != nil { + return nil, err + } + if l > len(data) { + return nil, errShortPacket + } + return data[:l], nil +} + type packetMarshaler interface { marshalPacket() (header, payload []byte, err error) } diff --git a/packet_test.go b/packet_test.go index d77b848a..a89e86ee 100644 --- a/packet_test.go +++ b/packet_test.go @@ -5,8 +5,10 @@ import ( "encoding" "errors" "io/ioutil" + "math" "os" "reflect" + "strconv" "testing" ) @@ -57,6 +59,229 @@ func TestUnmarshalFileStatExtendedOverflow(t *testing.T) { } } +func TestUnmarshalStatusShort(t *testing.T) { + for _, data := range [][]byte{ + nil, + {0x00}, + {0x00, 0x00, 0x00, 0x2a}, // only the id, no code + } { + if err := unmarshalStatus(42, data); err != errShortPacket { + t.Errorf("unmarshalStatus(%v): err = %v, want errShortPacket", data, err) + } + } +} + +func TestUnmarshalDataReply(t *testing.T) { + // id matches, declared length exceeds the bytes present. The 32-bit + // overflow of the length field itself is covered by TestUnmarshalCount. + b := marshalUint32(nil, 1) + b = marshalUint32(b, 0x7fffffff) + if _, err := unmarshalDataReply(1, b); err != errShortPacket { + t.Fatalf("expected errShortPacket, got %v", err) + } + + // truncated before the id (less than 4 bytes): rejected by unmarshalSID. + if _, err := unmarshalDataReply(1, []byte{0x00, 0x00}); err != errShortPacket { + t.Fatalf("expected errShortPacket, got %v", err) + } + + // truncated before the length field: id present, nothing after it. + if _, err := unmarshalDataReply(1, marshalUint32(nil, 1)); err != errShortPacket { + t.Fatalf("expected errShortPacket, got %v", err) + } + + // mismatched id. + b = marshalUint32(nil, 2) + b = marshalUint32(b, 0) + _, err := unmarshalDataReply(1, b) + var idErr *unexpectedIDErr + if !errors.As(err, &idErr) { + t.Fatalf("err = %v (%T), want *unexpectedIDErr", err, err) + } + if idErr.want != 1 { + t.Errorf("idErr.want = %d, want 1", idErr.want) + } + if idErr.got != 2 { + t.Errorf("idErr.got = %d, want 2", idErr.got) + } + + // well-formed reply: declared length shorter than the bytes present must + // return exactly the declared prefix, leaving any trailing bytes out. + payload := []byte("hello") + b = marshalUint32(nil, 1) + b = marshalUint32(b, uint32(len(payload))) + b = append(b, payload...) + b = append(b, "extra"...) // trailing bytes beyond the declared length + got, err := unmarshalDataReply(1, b) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !bytes.Equal(got, payload) { + t.Errorf("got %q, want %q", got, payload) + } +} + +// TestUnmarshalSID exercises the request-id check shared by every client reply +// decoder: a matching id returns the remaining bytes untouched, a mismatch +// reports both ids, and a truncated reply is rejected rather than panicking. +func TestUnmarshalSID(t *testing.T) { + // matching id: the bytes after the id must be returned verbatim. + b := marshalUint32(nil, 42) + b = append(b, 'x', 'y') + rest, err := unmarshalSID(42, b) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !bytes.Equal(rest, []byte("xy")) { + t.Errorf("rest = %q, want %q", rest, "xy") + } + + // matching id with no trailing bytes: rest is empty, not an error. + rest, err = unmarshalSID(7, marshalUint32(nil, 7)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(rest) != 0 { + t.Errorf("rest = % x, want empty", rest) + } + + // mismatched id: error carries both the expected and the received id. + _, err = unmarshalSID(1, marshalUint32(nil, 2)) + var idErr *unexpectedIDErr + if !errors.As(err, &idErr) { + t.Fatalf("err = %v (%T), want *unexpectedIDErr", err, err) + } + if idErr.want != 1 { + t.Errorf("idErr.want = %d, want 1", idErr.want) + } + if idErr.got != 2 { + t.Errorf("idErr.got = %d, want 2", idErr.got) + } + + // truncated before the id is fully present. + for _, b := range [][]byte{nil, {0x00}, {0x00, 0x00, 0x00}} { + if _, err := unmarshalSID(1, b); err != errShortPacket { + t.Errorf("unmarshalSID(%v): err = %v, want errShortPacket", b, err) + } + } +} + +// TestUnmarshalCount exercises the count decoder used for the SSH_FXP_NAME and +// SSH_FXP_DATA reply lengths. It returns the count and the remaining bytes, and +// guards the 32-bit overflow case that would otherwise produce a negative +// length in make/slice. +func TestUnmarshalCount(t *testing.T) { + // valid count: trailing bytes are returned untouched. + b := marshalUint32(nil, 3) + b = append(b, 'a', 'b', 'c') + count, rest, err := unmarshalCount(b) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if count != 3 { + t.Errorf("count = %d, want 3", count) + } + if !bytes.Equal(rest, []byte("abc")) { + t.Errorf("rest = %q, want %q", rest, "abc") + } + + // boundary: math.MaxInt32 (0x7fffffff) is the largest value that fits a + // non-negative int everywhere, so it must be accepted — this guards the + // off-by-one in the int(v) < 0 check. + count, _, err = unmarshalCount(marshalUint32(nil, 0x7fffffff)) + if err != nil { + t.Fatalf("MaxInt32: unexpected error: %v", err) + } + if count != math.MaxInt32 { + t.Errorf("MaxInt32: count = %d, want %d", count, math.MaxInt32) + } + + // truncated before the count is fully present. + for _, b := range [][]byte{nil, {0x00}, {0x00, 0x00, 0x00}} { + if _, _, err := unmarshalCount(b); err != errShortPacket { + t.Errorf("unmarshalCount(%v): err = %v, want errShortPacket", b, err) + } + } + + // A count larger than math.MaxInt32 does not fit in a 32-bit int. On those + // platforms it must be rejected with errLongPacket; on 64-bit platforms int + // is wide enough and the value is returned as-is. This check is always true + // on 64-bit and only ever trips on 32-bit. + big := marshalUint32(nil, 0x80000000) // math.MaxInt32 + 1, as a uint32 + count, _, err = unmarshalCount(big) + if strconv.IntSize == 32 { + if err != errLongPacket { + t.Errorf("32-bit: err = %v, want errLongPacket", err) + } + } else { + if err != nil { + t.Fatalf("64-bit: unexpected error: %v", err) + } + // int64 throughout: math.MaxInt32+1 overflows int on 32-bit at compile + // time, even though this branch only runs on 64-bit. + if int64(count) != int64(math.MaxInt32)+1 { + t.Errorf("64-bit: count = %d, want %d", count, int64(math.MaxInt32)+1) + } + } +} + +func TestUnmarshalStringSafe(t *testing.T) { + // well-formed: the declared bytes are returned as the string and any + // trailing bytes are handed back untouched as the remainder. + b := marshalString(nil, "hello") + b = append(b, 0xde, 0xad) + s, rest, err := unmarshalStringSafe(b) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if s != "hello" { + t.Errorf("s = %q, want %q", s, "hello") + } + if !bytes.Equal(rest, []byte{0xde, 0xad}) { + t.Errorf("rest = % x, want de ad", rest) + } + + // empty string: a zero length yields "" and leaves the remainder untouched. + empty := append(marshalUint32(nil, 0), 'x', 'y') + s, rest, err = unmarshalStringSafe(empty) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if s != "" { + t.Errorf("s = %q, want empty", s) + } + if !bytes.Equal(rest, []byte("xy")) { + t.Errorf("rest = %q, want %q", rest, "xy") + } + + // exact length: the declared bytes consume the whole slice, empty remainder. + s, rest, err = unmarshalStringSafe(marshalString(nil, "abc")) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if s != "abc" { + t.Errorf("s = %q, want %q", s, "abc") + } + if len(rest) != 0 { + t.Errorf("rest = % x, want empty", rest) + } + + // truncated before the length prefix is fully present. + for _, b := range [][]byte{nil, {0x00}, {0x00, 0x00, 0x00}} { + if _, _, err := unmarshalStringSafe(b); err != errShortPacket { + t.Errorf("unmarshalStringSafe(% x): err = %v, want errShortPacket", b, err) + } + } + + // declared length exceeds the bytes actually present. The 32-bit overflow + // of the length field itself is covered by TestUnmarshalCount. + short := marshalUint32(nil, 5) + short = append(short, 'a', 'b') // only 2 of the 5 declared bytes + if _, _, err := unmarshalStringSafe(short); err != errShortPacket { + t.Errorf("unmarshalStringSafe(short): err = %v, want errShortPacket", err) + } +} + func TestMarshalUint64(t *testing.T) { var tests = []struct { v uint64