624 lines
17 KiB
Go
624 lines
17 KiB
Go
//go:build linux || freebsd
|
|
// +build linux freebsd
|
|
|
|
package sftp
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"crypto/rsa"
|
|
"encoding/binary"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/url"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
|
|
"codeberg.org/petrbalvin/goget/internal/core"
|
|
"golang.org/x/crypto/ssh"
|
|
)
|
|
|
|
// startTestServer starts a minimal SSH+SFTP server for testing.
|
|
func startTestServer(t *testing.T) (addr string, cleanup func()) {
|
|
t.Helper()
|
|
|
|
config := &ssh.ServerConfig{
|
|
NoClientAuth: true,
|
|
}
|
|
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
|
|
if err != nil {
|
|
t.Fatalf("generate key: %v", err)
|
|
}
|
|
signer, err := ssh.NewSignerFromKey(privateKey)
|
|
if err != nil {
|
|
t.Fatalf("new signer: %v", err)
|
|
}
|
|
config.AddHostKey(signer)
|
|
|
|
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
|
if err != nil {
|
|
t.Fatalf("listen: %v", err)
|
|
}
|
|
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
// Accept connections in a loop so the parallel-download tests
|
|
// (which open N SFTP sessions) can run. Each connection is
|
|
// served in its own goroutine, so the N pool connections can be
|
|
// alive concurrently as long as the test goroutine keeps
|
|
// accepting.
|
|
for {
|
|
conn, err := listener.Accept()
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
go func(conn net.Conn) {
|
|
defer conn.Close()
|
|
sshConn, chans, reqs, err := ssh.NewServerConn(conn, config)
|
|
if err != nil {
|
|
return
|
|
}
|
|
defer sshConn.Close()
|
|
go ssh.DiscardRequests(reqs)
|
|
|
|
for newChannel := range chans {
|
|
if newChannel.ChannelType() != "session" {
|
|
newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
|
|
continue
|
|
}
|
|
channel, requests, err := newChannel.Accept()
|
|
if err != nil {
|
|
return
|
|
}
|
|
go func(ch ssh.Channel, reqs <-chan *ssh.Request) {
|
|
for req := range reqs {
|
|
ok := false
|
|
if req.Type == "subsystem" && len(req.Payload) > 4 {
|
|
l := binary.BigEndian.Uint32(req.Payload[:4])
|
|
name := string(req.Payload[4 : 4+l])
|
|
if name == "sftp" {
|
|
ok = true
|
|
if req.WantReply {
|
|
req.Reply(ok, nil)
|
|
}
|
|
handleSFTPSession(ch)
|
|
return
|
|
}
|
|
}
|
|
if req.WantReply {
|
|
req.Reply(ok, nil)
|
|
}
|
|
}
|
|
ch.Close()
|
|
}(channel, requests)
|
|
}
|
|
}(conn)
|
|
}
|
|
}()
|
|
|
|
return listener.Addr().String(), func() {
|
|
listener.Close()
|
|
<-done
|
|
}
|
|
}
|
|
|
|
// handleSFTPSession handles a minimal SFTP v3 session.
|
|
func handleSFTPSession(ch ssh.Channel) {
|
|
defer ch.Close()
|
|
|
|
// Expect INIT
|
|
pkt, err := readTestPacket(ch)
|
|
if err != nil || len(pkt) < 5 || pkt[0] != sshFxpInit {
|
|
return
|
|
}
|
|
// Respond VERSION
|
|
var resp bytes.Buffer
|
|
resp.WriteByte(sshFxpVersion)
|
|
writeUint32(&resp, 3)
|
|
writeTestPacket(ch, resp.Bytes())
|
|
|
|
handles := make(map[string]string)
|
|
dirHandles := make(map[string][]string)
|
|
dirHandlePaths := make(map[string]string)
|
|
handleCounter := 0
|
|
|
|
for {
|
|
pkt, err := readTestPacket(ch)
|
|
if err != nil {
|
|
return
|
|
}
|
|
if len(pkt) < 1 {
|
|
continue
|
|
}
|
|
reqID, _ := readUint32(pkt, 1)
|
|
|
|
switch pkt[0] {
|
|
case sshFxpRealpath:
|
|
path, _ := readString(pkt, 5)
|
|
if path == "." {
|
|
path = "/home/test"
|
|
}
|
|
var r bytes.Buffer
|
|
r.WriteByte(sshFxpName)
|
|
writeUint32(&r, reqID)
|
|
writeUint32(&r, 1)
|
|
writeString(&r, path)
|
|
writeString(&r, path)
|
|
writeUint32(&r, 0) // attrs
|
|
writeTestPacket(ch, r.Bytes())
|
|
|
|
case sshFxpStat:
|
|
path, _ := readString(pkt, 5)
|
|
var r bytes.Buffer
|
|
r.WriteByte(sshFxpAttrs)
|
|
writeUint32(&r, reqID)
|
|
fs := testFSLookup(path)
|
|
if fs != nil && fs.isDir {
|
|
// mode flag only
|
|
writeUint32(&r, 0x4)
|
|
writeUint32(&r, 0x4000) // S_IFDIR
|
|
writeTestPacket(ch, r.Bytes())
|
|
continue
|
|
}
|
|
if fs != nil && !fs.isDir {
|
|
// size flag only
|
|
writeUint32(&r, 0x1)
|
|
writeUint64(&r, uint64(len(fs.content)))
|
|
writeTestPacket(ch, r.Bytes())
|
|
continue
|
|
}
|
|
// legacy single-file mock behaviour
|
|
writeUint32(&r, 0x1) // size flag
|
|
size := uint64(len(testFileContent))
|
|
if testOverrideAttrSize != nil {
|
|
size = *testOverrideAttrSize
|
|
}
|
|
writeUint64(&r, size)
|
|
writeTestPacket(ch, r.Bytes())
|
|
|
|
case sshFxpOpendir:
|
|
path, _ := readString(pkt, 5)
|
|
fs := testFSLookup(path)
|
|
if fs == nil || !fs.isDir {
|
|
status(ch, reqID, sshFxNoSuchFile, "no such directory")
|
|
continue
|
|
}
|
|
handleCounter++
|
|
h := fmt.Sprintf("dhandle%d", handleCounter)
|
|
dirHandles[h] = fs.children
|
|
dirHandlePaths[h] = path
|
|
var r bytes.Buffer
|
|
r.WriteByte(sshFxpHandle)
|
|
writeUint32(&r, reqID)
|
|
writeString(&r, h)
|
|
writeTestPacket(ch, r.Bytes())
|
|
|
|
case sshFxpReaddir:
|
|
handle, _ := readString(pkt, 5)
|
|
children, ok := dirHandles[handle]
|
|
if !ok {
|
|
status(ch, reqID, sshFxFailure, "no such dir handle")
|
|
continue
|
|
}
|
|
parentPath := dirHandlePaths[handle]
|
|
var r bytes.Buffer
|
|
r.WriteByte(sshFxpName)
|
|
writeUint32(&r, reqID)
|
|
writeUint32(&r, uint32(len(children)))
|
|
for _, name := range children {
|
|
writeString(&r, name)
|
|
writeString(&r, name)
|
|
childPath := parentPath + "/" + name
|
|
childFS := testFSLookup(childPath)
|
|
isDir := childFS != nil && childFS.isDir
|
|
// attrs: size flag (0x1) and mode flag (0x4) so the
|
|
// client can both classify the entry (dir vs file) and
|
|
// read the file size for parallel accounting.
|
|
if isDir {
|
|
writeUint32(&r, 0x4)
|
|
writeUint32(&r, 0x4000) // S_IFDIR
|
|
} else {
|
|
writeUint32(&r, 0x5) // size + mode
|
|
var sz uint64
|
|
if childFS != nil {
|
|
sz = uint64(len(childFS.content))
|
|
} else {
|
|
sz = uint64(len(testFileContent))
|
|
}
|
|
writeUint64(&r, sz)
|
|
writeUint32(&r, 0x8000) // S_IFREG
|
|
}
|
|
}
|
|
writeTestPacket(ch, r.Bytes())
|
|
|
|
case sshFxpOpen:
|
|
path, off := readString(pkt, 5)
|
|
pflags, _ := readUint32(pkt, off)
|
|
handleCounter++
|
|
h := fmt.Sprintf("handle%d", handleCounter)
|
|
handles[h] = path
|
|
var r bytes.Buffer
|
|
if pflags&sshFxfRead != 0 {
|
|
r.WriteByte(sshFxpHandle)
|
|
writeUint32(&r, reqID)
|
|
writeString(&r, h)
|
|
writeTestPacket(ch, r.Bytes())
|
|
} else if pflags&sshFxfWrite != 0 {
|
|
r.WriteByte(sshFxpHandle)
|
|
writeUint32(&r, reqID)
|
|
writeString(&r, h)
|
|
writeTestPacket(ch, r.Bytes())
|
|
} else {
|
|
status(ch, reqID, sshFxFailure, "unsupported flags")
|
|
}
|
|
|
|
case sshFxpRead:
|
|
handle, off := readString(pkt, 5)
|
|
offset, off := readUint64(pkt, off)
|
|
length, _ := readUint32(pkt, off)
|
|
_ = length
|
|
path := handles[handle]
|
|
if path == "" {
|
|
status(ch, reqID, sshFxNoSuchFile, "no such handle")
|
|
continue
|
|
}
|
|
var data []byte
|
|
if fs := testFSLookup(path); fs != nil && !fs.isDir {
|
|
data = fs.content
|
|
} else {
|
|
data = testFileContent
|
|
}
|
|
if int(offset) >= len(data) {
|
|
status(ch, reqID, sshFxEOF, "")
|
|
continue
|
|
}
|
|
end := int(offset) + int(length)
|
|
if end > len(data) {
|
|
end = len(data)
|
|
}
|
|
chunk := data[offset:end]
|
|
var r bytes.Buffer
|
|
r.WriteByte(sshFxpData)
|
|
writeUint32(&r, reqID)
|
|
writeUint32(&r, uint32(len(chunk)))
|
|
r.Write(chunk)
|
|
writeTestPacket(ch, r.Bytes())
|
|
|
|
case sshFxpClose:
|
|
handle, _ := readString(pkt, 5)
|
|
delete(handles, handle)
|
|
delete(dirHandles, handle)
|
|
var r bytes.Buffer
|
|
r.WriteByte(sshFxpStatus)
|
|
writeUint32(&r, reqID)
|
|
writeUint32(&r, sshFxOk)
|
|
writeString(&r, "")
|
|
writeString(&r, "")
|
|
writeTestPacket(ch, r.Bytes())
|
|
|
|
default:
|
|
status(ch, reqID, sshFxOpUnsupported, "unsupported")
|
|
}
|
|
}
|
|
}
|
|
|
|
// testFSLookup is a tiny helper that hides the nil-map check so the
|
|
// mock switch arms above can stay readable.
|
|
func testFSLookup(path string) *testFSEntry {
|
|
if testFS == nil {
|
|
return nil
|
|
}
|
|
return testFS[path]
|
|
}
|
|
|
|
func status(ch ssh.Channel, reqID, code uint32, msg string) {
|
|
var r bytes.Buffer
|
|
r.WriteByte(sshFxpStatus)
|
|
writeUint32(&r, reqID)
|
|
writeUint32(&r, code)
|
|
writeString(&r, msg)
|
|
writeString(&r, "")
|
|
writeTestPacket(ch, r.Bytes())
|
|
}
|
|
|
|
func readTestPacket(r io.Reader) ([]byte, error) {
|
|
var lenBuf [4]byte
|
|
if _, err := io.ReadFull(r, lenBuf[:]); err != nil {
|
|
return nil, err
|
|
}
|
|
length := binary.BigEndian.Uint32(lenBuf[:])
|
|
if length > 256*1024 {
|
|
return nil, fmt.Errorf("packet too large")
|
|
}
|
|
pkt := make([]byte, length)
|
|
if _, err := io.ReadFull(r, pkt); err != nil {
|
|
return nil, err
|
|
}
|
|
return pkt, nil
|
|
}
|
|
|
|
func writeTestPacket(w io.Writer, pkt []byte) error {
|
|
var lenBuf [4]byte
|
|
binary.BigEndian.PutUint32(lenBuf[:], uint32(len(pkt)))
|
|
if _, err := w.Write(lenBuf[:]); err != nil {
|
|
return err
|
|
}
|
|
_, err := w.Write(pkt)
|
|
return err
|
|
}
|
|
|
|
var testFileContent = []byte("hello sftp world")
|
|
|
|
// testOverrideAttrSize, when non-nil, replaces the value reported by the
|
|
// mock server's sshFxpStat handler. Tests use it to simulate a server
|
|
// that returns a wrong (typically too large) size attribute so we can
|
|
// verify the client does not blindly trust attr.Size for
|
|
// BytesDownloaded accounting.
|
|
var testOverrideAttrSize *uint64
|
|
|
|
// testFSEntry describes a single in-memory file or directory used by
|
|
// the recursive-download tests. isDir==true entries carry a children
|
|
// slice; isDir==false entries carry the file content.
|
|
type testFSEntry struct {
|
|
isDir bool
|
|
content []byte
|
|
children []string
|
|
}
|
|
|
|
// testFS is an optional in-memory filesystem consulted by the mock
|
|
// server before falling back to the default testFileContent behaviour.
|
|
// Tests that exercise recursive downloads set it up and clear it via
|
|
// t.Cleanup. Keys are absolute paths; paths NOT present in the map
|
|
// keep the legacy single-file mock semantics so existing
|
|
// non-recursive tests continue to pass.
|
|
var testFS map[string]*testFSEntry
|
|
|
|
func TestDownload(t *testing.T) {
|
|
addr, cleanup := startTestServer(t)
|
|
defer cleanup()
|
|
|
|
u, _ := url.Parse("sftp://user@" + addr + "/testfile")
|
|
outDir := t.TempDir()
|
|
outputPath := filepath.Join(outDir, "out")
|
|
|
|
req := &core.DownloadRequest{
|
|
URL: u,
|
|
Output: outputPath,
|
|
Ctx: context.Background(),
|
|
SSHInsecure: true,
|
|
}
|
|
|
|
result, err := download(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("download failed: %v", err)
|
|
}
|
|
|
|
data, err := os.ReadFile(outputPath)
|
|
if err != nil {
|
|
t.Fatalf("read output: %v", err)
|
|
}
|
|
if !bytes.Equal(data, testFileContent) {
|
|
t.Errorf("got %q, want %q", data, testFileContent)
|
|
}
|
|
if result.BytesDownloaded != int64(len(testFileContent)) {
|
|
t.Errorf("bytes=%d, want %d", result.BytesDownloaded, len(testFileContent))
|
|
}
|
|
if result.Protocol != "sftp" {
|
|
t.Errorf("protocol=%q, want sftp", result.Protocol)
|
|
}
|
|
}
|
|
|
|
// TestDownloadWithStaleAttrSize is a regression guard for the BACKLOG
|
|
// entry "SFTP BytesDownloaded unconditionally overwritten". The previous
|
|
// implementation set result.BytesDownloaded = int64(attr.Size) when
|
|
// attr.Size > 0, ignoring the locally-tracked `total` counter. If the
|
|
// server reported a stale (typically too large) size, the result would
|
|
// over-report the bytes actually transferred — which is exactly what
|
|
// this test exercises.
|
|
func TestDownloadWithStaleAttrSize(t *testing.T) {
|
|
// Simulate a server that reports 10 KiB even though the file is
|
|
// only len(testFileContent) bytes long. After the fix, the result
|
|
// must reflect the actual bytes received, not the stale attribute.
|
|
wrong := uint64(10 * 1024)
|
|
testOverrideAttrSize = &wrong
|
|
t.Cleanup(func() { testOverrideAttrSize = nil })
|
|
|
|
addr, cleanup := startTestServer(t)
|
|
defer cleanup()
|
|
|
|
u, _ := url.Parse("sftp://user@" + addr + "/testfile")
|
|
outDir := t.TempDir()
|
|
outputPath := filepath.Join(outDir, "out")
|
|
|
|
req := &core.DownloadRequest{
|
|
URL: u,
|
|
Output: outputPath,
|
|
Ctx: context.Background(),
|
|
SSHInsecure: true,
|
|
}
|
|
|
|
result, err := download(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("download failed: %v", err)
|
|
}
|
|
|
|
// The downloaded file must still hold the real (small) payload.
|
|
data, err := os.ReadFile(outputPath)
|
|
if err != nil {
|
|
t.Fatalf("read output: %v", err)
|
|
}
|
|
if !bytes.Equal(data, testFileContent) {
|
|
t.Errorf("got %q, want %q", data, testFileContent)
|
|
}
|
|
|
|
// And BytesDownloaded must reflect the actual bytes received, not
|
|
// the stale attr.Size the server returned.
|
|
if result.BytesDownloaded != int64(len(testFileContent)) {
|
|
t.Errorf("BytesDownloaded = %d, want %d (real payload size, not stale attr.Size=%d)",
|
|
result.BytesDownloaded, len(testFileContent), wrong)
|
|
}
|
|
}
|
|
|
|
func TestParseHeader(t *testing.T) {
|
|
var buf bytes.Buffer
|
|
writeUint32(&buf, 0x05) // size + mode flags
|
|
writeUint64(&buf, 60)
|
|
writeUint32(&buf, 0644)
|
|
|
|
attr, off := readAttrs(buf.Bytes(), 0)
|
|
if attr.Size != 60 {
|
|
t.Errorf("size=%d, want 60", attr.Size)
|
|
}
|
|
if off != 16 {
|
|
t.Errorf("off=%d, want 16", off)
|
|
}
|
|
}
|
|
|
|
func TestIsEOF(t *testing.T) {
|
|
if !isEOF(fmt.Errorf("sftp status %d", sshFxEOF)) {
|
|
t.Error("expected true")
|
|
}
|
|
if isEOF(fmt.Errorf("sftp status %d", sshFxFailure)) {
|
|
t.Error("expected false")
|
|
}
|
|
}
|
|
|
|
func TestDialTimeout(t *testing.T) {
|
|
u, _ := url.Parse("sftp://user@127.0.0.1:1/")
|
|
_, err := dial(context.Background(), u, "", true, "")
|
|
if err == nil {
|
|
t.Fatal("expected error")
|
|
}
|
|
}
|
|
|
|
func TestSFTPRecursiveDownloadParallel(t *testing.T) {
|
|
addr, cleanup := startTestServer(t)
|
|
defer cleanup()
|
|
|
|
setupRecursiveFS := func(t *testing.T) {
|
|
t.Helper()
|
|
testFS = map[string]*testFSEntry{
|
|
"/home/test/dir": {isDir: true, children: []string{"a.txt", "b.txt", "c.txt", "sub"}},
|
|
"/home/test/dir/a.txt": {isDir: false, content: []byte("alpha")},
|
|
"/home/test/dir/b.txt": {isDir: false, content: []byte("bravo-bravo")},
|
|
"/home/test/dir/c.txt": {isDir: false, content: []byte("charlie")},
|
|
"/home/test/dir/sub": {isDir: true, children: []string{"x.txt", "y.txt"}},
|
|
"/home/test/dir/sub/x.txt": {isDir: false, content: []byte("x-ray")},
|
|
"/home/test/dir/sub/y.txt": {isDir: false, content: []byte("yankee")},
|
|
}
|
|
t.Cleanup(func() { testFS = nil })
|
|
}
|
|
|
|
t.Run("parallel=3 downloads all files", func(t *testing.T) {
|
|
setupRecursiveFS(t)
|
|
u, _ := url.Parse("sftp://user@" + addr + "/home/test/dir")
|
|
outDir := t.TempDir()
|
|
req := &core.DownloadRequest{
|
|
URL: u,
|
|
Output: outDir,
|
|
Recursive: true,
|
|
RecursiveParallel: 3,
|
|
SSHInsecure: true,
|
|
}
|
|
result, err := download(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("recursive parallel download failed: %v", err)
|
|
}
|
|
if result.ChunksCount != 5 {
|
|
t.Errorf("ChunksCount = %d, want 5 (3 top-level files + 2 in subdir)", result.ChunksCount)
|
|
}
|
|
// 5 + 11 + 7 + 5 + 6 = 34 bytes total content
|
|
const wantBytes = int64(5 + 11 + 7 + 5 + 6)
|
|
if result.BytesDownloaded != wantBytes {
|
|
t.Errorf("BytesDownloaded = %d, want %d", result.BytesDownloaded, wantBytes)
|
|
}
|
|
// Verify all 5 files actually exist with correct content.
|
|
wantFiles := map[string]string{
|
|
"a.txt": "alpha",
|
|
"b.txt": "bravo-bravo",
|
|
"c.txt": "charlie",
|
|
"sub/x.txt": "x-ray",
|
|
"sub/y.txt": "yankee",
|
|
}
|
|
for name, want := range wantFiles {
|
|
got, err := os.ReadFile(filepath.Join(outDir, name))
|
|
if err != nil {
|
|
t.Errorf("file %s: %v", name, err)
|
|
continue
|
|
}
|
|
if string(got) != want {
|
|
t.Errorf("file %s = %q, want %q", name, string(got), want)
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("parallel=0 falls back to sequential", func(t *testing.T) {
|
|
setupRecursiveFS(t)
|
|
u, _ := url.Parse("sftp://user@" + addr + "/home/test/dir")
|
|
outDir := t.TempDir()
|
|
req := &core.DownloadRequest{
|
|
URL: u,
|
|
Output: outDir,
|
|
Recursive: true,
|
|
RecursiveParallel: 0,
|
|
SSHInsecure: true,
|
|
}
|
|
result, err := download(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("recursive sequential download failed: %v", err)
|
|
}
|
|
if result.ChunksCount != 5 {
|
|
t.Errorf("ChunksCount = %d, want 5", result.ChunksCount)
|
|
}
|
|
})
|
|
|
|
t.Run("parallel=1 falls back to sequential", func(t *testing.T) {
|
|
setupRecursiveFS(t)
|
|
u, _ := url.Parse("sftp://user@" + addr + "/home/test/dir")
|
|
outDir := t.TempDir()
|
|
req := &core.DownloadRequest{
|
|
URL: u,
|
|
Output: outDir,
|
|
Recursive: true,
|
|
RecursiveParallel: 1,
|
|
SSHInsecure: true,
|
|
}
|
|
result, err := download(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("download with parallel=1 failed: %v", err)
|
|
}
|
|
if result.ChunksCount != 5 {
|
|
t.Errorf("ChunksCount = %d, want 5", result.ChunksCount)
|
|
}
|
|
})
|
|
|
|
t.Run("subdir propagation with parallel", func(t *testing.T) {
|
|
setupRecursiveFS(t)
|
|
u, _ := url.Parse("sftp://user@" + addr + "/home/test/dir")
|
|
outDir := t.TempDir()
|
|
req := &core.DownloadRequest{
|
|
URL: u,
|
|
Output: outDir,
|
|
Recursive: true,
|
|
RecursiveParallel: 4,
|
|
SSHInsecure: true,
|
|
}
|
|
_, err := download(context.Background(), req)
|
|
if err != nil {
|
|
t.Fatalf("download failed: %v", err)
|
|
}
|
|
// Both subdir files should be present and correct.
|
|
for _, name := range []string{"sub/x.txt", "sub/y.txt"} {
|
|
path := filepath.Join(outDir, name)
|
|
if _, err := os.Stat(path); err != nil {
|
|
t.Errorf("expected file %s to exist: %v", name, err)
|
|
}
|
|
}
|
|
})
|
|
}
|