feat: WithSubsystem (#224)

* feat: better integration with pkg/sftp This would allow users to more easily provide both SCP and SFTP servers to their users. closes #40 Signed-off-by: Carlos Alexandro Becker <[email protected]> * test: fix Signed-off-by: Carlos Alexandro Becker <[email protected]> * test: add tests Signed-off-by: Carlos Alexandro Becker <[email protected]> * feat: sftp refactory * fix: aymans suggestions Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: aymans suggestions Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: make sftp an example instead Signed-off-by: Carlos Alexandro Becker <[email protected]> * chore: update docs * fix: tests Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: unexport Signed-off-by: Carlos Alexandro Becker <[email protected]> * fix: unexport Signed-off-by: Carlos Alexandro Becker <[email protected]> --------- Signed-off-by: Carlos Alexandro Becker <[email protected]>
charmbracelet · Jan 22, 2024 · e5d20f5 · e5d20f5
1 parent 1fbb5ec
commit e5d20f5
Show file tree

Hide file tree

Showing 7 changed files with 154 additions and 3 deletions.
diff --git a/examples/go.mod b/examples/go.mod
@@ -9,6 +9,8 @@ require (
  github.com/charmbracelet/log v0.3.1
  github.com/charmbracelet/ssh v0.0.0-20240118173142-6d7cf11c8371
  github.com/charmbracelet/wish v0.5.0
+ github.com/muesli/termenv v0.15.2
+ github.com/pkg/sftp v1.13.6
  github.com/spf13/cobra v1.8.0
  golang.org/x/crypto v0.18.0
 )
@@ -36,14 +38,14 @@ require (
  github.com/inconshreveable/mousetrap v1.1.0 // indirect
  github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect
  github.com/kevinburke/ssh_config v1.2.0 // indirect
+ github.com/kr/fs v0.1.0 // indirect
  github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
  github.com/mattn/go-isatty v0.0.18 // indirect
  github.com/mattn/go-localereader v0.0.1 // indirect
  github.com/mattn/go-runewidth v0.0.15 // indirect
  github.com/muesli/ansi v0.0.0-20211018074035-2e021307bc4b // indirect
  github.com/muesli/cancelreader v0.2.2 // indirect
  github.com/muesli/reflow v0.3.0 // indirect
- github.com/muesli/termenv v0.15.2 // indirect
  github.com/pjbgf/sha1cd v0.3.0 // indirect
  github.com/rivo/uniseg v0.2.0 // indirect
  github.com/sergi/go-diff v1.1.0 // indirect

diff --git a/examples/go.sum b/examples/go.sum
@@ -64,6 +64,8 @@ github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 h1:BQSFePA1RWJOl
 github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99/go.mod h1:1lJo3i6rXxKeerYnT8Nvf0QmHCRC1n8sfWVwXF2Frvo=
 github.com/kevinburke/ssh_config v1.2.0 h1:x584FjTGwHzMwvHx18PXxbBVzfnxogHaAReU4gf13a4=
 github.com/kevinburke/ssh_config v1.2.0/go.mod h1:CT57kijsi8u/K/BOFA39wgDQJ9CxiF4nAY/ojJ6r6mM=
+github.com/kr/fs v0.1.0 h1:Jskdu9ieNAYnjxsi0LbQp1ulIKZV1LAFgK1tWhpZgl8=
+github.com/kr/fs v0.1.0/go.mod h1:FFnZGqtBN9Gxj7eW1uZ42v5BccTP0vu6NEaFoC2HwRg=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
@@ -92,6 +94,8 @@ github.com/pjbgf/sha1cd v0.3.0 h1:4D5XXmUUBUl/xQ6IjCkEAbqXskkq/4O7LmGn0AqMDs4=
 github.com/pjbgf/sha1cd v0.3.0/go.mod h1:nZ1rrWOcGJ5uZgEEVL1VUM9iRQiZvWdbZjkKyFzPPsI=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
+github.com/pkg/sftp v1.13.6 h1:JFZT4XbOU7l77xGSpOdW+pwIMqP044IyjXX6FGyEKFo=
+github.com/pkg/sftp v1.13.6/go.mod h1:tz1ryNURKu77RL+GuCzmoJYxQczL3wLNNpPWagdg4Qk=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
@@ -109,8 +113,11 @@ github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyh
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
+github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
 github.com/u-root/gobusybox/src v0.0.0-20221229083637-46b2883a7f90 h1:zTk5683I9K62wtZ6eUa6vu6IWwVHXPnoKK5n2unAwv0=
 github.com/u-root/u-root v0.11.0 h1:6gCZLOeRyevw7gbTwMj3fKxnr9+yHFlgF3N7udUVNO8=
@@ -121,6 +128,7 @@ github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5t
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
+golang.org/x/crypto v0.1.0/go.mod h1:RecgLatLF4+eUMCP1PoPZQb+cVrJcOPbHkTkbkB9sbw=
 golang.org/x/crypto v0.3.1-0.20221117191849-2c476679df9a/go.mod h1:hebNnKkNXi2UzZN1eVRvBB7co0a+JxK6XbPiWVs/3J4=
 golang.org/x/crypto v0.7.0/go.mod h1:pYwdfH91IfpZVANVyUOhSIPZaFoJGxTFbZhFTx+dXZU=
 golang.org/x/crypto v0.18.0 h1:PGVlW0xEltQnzFZ55hkuX5+KLyrMYhHld1YHO4AKcdc=
@@ -135,6 +143,7 @@ golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLL
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
+golang.org/x/net v0.1.0/go.mod h1:Cx3nUiGt4eDBEyega/BKRp+/AlGL8hYe7U9odMt2Cco=
 golang.org/x/net v0.2.0/go.mod h1:KqCZLdyyvdV855qA2rE3GC2aiw5xGR5TEjj8smXukLY=
 golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
 golang.org/x/net v0.8.0/go.mod h1:QVkue5JL9kW//ek3r6jTKnTFis1tRmNAW2P1shuFdJc=
@@ -163,6 +172,7 @@ golang.org/x/sys v0.16.0 h1:xWw16ngr6ZMtmxDyKyIgsE93KNKz5HKmMa3b8ALHidU=
 golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
+golang.org/x/term v0.1.0/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/term v0.2.0/go.mod h1:TVmDHMZPmdnySmBfhjOoOdhjzdE1h4u1VwSiw2l1Nuc=
 golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
 golang.org/x/term v0.6.0/go.mod h1:m6U89DPEgQRMq3DNkDClhWw02AUbt2daBVO4cn4Hv9U=
@@ -191,5 +201,6 @@ gopkg.in/warnings.v0 v0.1.2 h1:wFXVbFY8DY5/xOe1ECiWdKCzZlxgshcYVNkBHstARME=
 gopkg.in/warnings.v0 v0.1.2/go.mod h1:jksf8JmL6Qr/oQM2OXTHunEvvTAsrWBLb6OOjuVWRNI=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/examples/scp/main.go b/examples/scp/main.go
@@ -6,15 +6,19 @@ import (
  "context"
  "errors"
  "fmt"
+ "io"
+ "io/fs"
  "os"
  "os/signal"
+ "path/filepath"
  "syscall"
  "time"
 
  "github.com/charmbracelet/log"
  "github.com/charmbracelet/ssh"
  "github.com/charmbracelet/wish"
  "github.com/charmbracelet/wish/scp"
+ "github.com/pkg/sftp"
 )
 
 const (
@@ -23,10 +27,12 @@ const (
 )
 
 func main() {
- handler := scp.NewFileSystemHandler("./examples/scp/testdata")
+ root, _ := filepath.Abs("./examples/scp/testdata")
+ handler := scp.NewFileSystemHandler(root)
  s, err := wish.NewServer(
  wish.WithAddress(fmt.Sprintf("%s:%d", host, port)),
  wish.WithHostKeyPath(".ssh/term_info_ed25519"),
+ wish.WithSubsystem("sftp", sftpSubsystem(root)),
  wish.WithMiddleware(
  scp.Middleware(handler, handler),
  ),
@@ -53,3 +59,109 @@ func main() {
  log.Error("could not stop server", "error", err)
  }
 }
+
+func sftpSubsystem(root string) ssh.SubsystemHandler {
+ return func(s ssh.Session) {
+ log.Info("sftp", "root", root)
+ fs := &sftpHandler{root}
+ srv := sftp.NewRequestServer(s, sftp.Handlers{
+ FileList: fs,
+ FileGet: fs,
+ })
+ if err := srv.Serve(); err == io.EOF {
+ if err := srv.Close(); err != nil {
+ wish.Fatalln(s, "sftp:", err)
+ }
+ } else if err != nil {
+ wish.Fatalln(s, "sftp:", err)
+ }
+ }
+}
+
+// Example readonly handler implementation for sftp.
+//
+// other example implementations:
+// - https:/gravitational/teleport/blob/f57dc2fe2a9900ec198779aae747ac4f833b278d/tool/teleport/common/sftp.go
+// - https:/minio/minio/blob/c66c5828eacb4a7fa9a49b4c890c77dd8684b171/cmd/sftp-server.go
+type sftpHandler struct {
+ root string
+}
+
+var (
+ _ sftp.FileLister = &sftpHandler{}
+ _ sftp.FileReader = &sftpHandler{}
+)
+
+type listerAt []fs.FileInfo
+
+func (l listerAt) ListAt(ls []fs.FileInfo, offset int64) (int, error) {
+ if offset >= int64(len(l)) {
+ return 0, io.EOF
+ }
+ n := copy(ls, l[offset:])
+ if n < len(ls) {
+ return n, io.EOF
+ }
+ return n, nil
+}
+
+// Fileread implements sftp.FileReader.
+func (s *sftpHandler) Fileread(r *sftp.Request) (io.ReaderAt, error) {
+ var flags int
+ pflags := r.Pflags()
+ if pflags.Append {
+ flags |= os.O_APPEND
+ }
+ if pflags.Creat {
+ flags |= os.O_CREATE
+ }
+ if pflags.Excl {
+ flags |= os.O_EXCL
+ }
+ if pflags.Trunc {
+ flags |= os.O_TRUNC
+ }
+
+ if pflags.Read && pflags.Write {
+ flags |= os.O_RDWR
+ } else if pflags.Read {
+ flags |= os.O_RDONLY
+ } else if pflags.Write {
+ flags |= os.O_WRONLY
+ }
+
+ f, err := os.OpenFile(filepath.Join(s.root, r.Filepath), flags, 0600)
+ if err != nil {
+ return nil, err
+ }
+
+ return f, nil
+}
+
+// Filelist implements sftp.FileLister.
+func (s *sftpHandler) Filelist(r *sftp.Request) (sftp.ListerAt, error) {
+ switch r.Method {
+ case "List":
+ entries, err := os.ReadDir(filepath.Join(s.root, r.Filepath))
+ if err != nil {
+ return nil, fmt.Errorf("sftp: %w", err)
+ }
+ infos := make([]fs.FileInfo, len(entries))
+ for i, entry := range entries {
+ info, err := entry.Info()
+ if err != nil {
+ return nil, err
+ }
+ infos[i] = info
+ }
+ return listerAt(infos), nil
+ case "Stat":
+ fi, err := os.Stat(filepath.Join(s.root, r.Filepath))
+ if err != nil {
+ return nil, err
+ }
+ return listerAt{fi}, nil
+ default:
+ return nil, sftp.ErrSSHFxOpUnsupported
+ }
+}
diff --git a/options.go b/options.go
@@ -199,3 +199,15 @@ func WithMaxTimeout(d time.Duration) ssh.Option {
  return nil
  }
 }
+
+// WithSubsystem returns an ssh.Option that sets the subsystem
+// handler for a given protocol.
+func WithSubsystem(key string, h ssh.SubsystemHandler) ssh.Option {
+ return func(s *ssh.Server) error {
+ if s.SubsystemHandlers == nil {
+ s.SubsystemHandlers = map[string]ssh.SubsystemHandler{}
+ }
+ s.SubsystemHandlers[key] = h
+ return nil
+ }
+}
diff --git a/options_test.go b/options_test.go
@@ -13,6 +13,19 @@ import (
  gossh "golang.org/x/crypto/ssh"
 )
 
+func TestWithSubsystem(t *testing.T) {
+ srv := &ssh.Server{
+ Handler: func(s ssh.Session) {},
+ }
+ requireNoError(t, WithSubsystem("foo", func(s ssh.Session) {})(srv))
+ if srv.SubsystemHandlers == nil {
+ t.Fatalf("should not have been nil")
+ }
+ if _, ok := srv.SubsystemHandlers["foo"]; !ok {
+ t.Fatalf("should have set the foo subsystem handler")
+ }
+}
+
 func TestWithBanner(t *testing.T) {
  const banner = "a banner"
  var got string

diff --git a/scp/filesystem.go b/scp/filesystem.go
@@ -12,6 +12,7 @@ import (
  "github.com/charmbracelet/ssh"
 )
 
+// fileSystemHandler is a Handler implementation for a given root path.
 type fileSystemHandler struct{ root string }
 
 var _ Handler = &fileSystemHandler{}

diff --git a/scp/filesystem_test.go b/scp/filesystem_test.go
@@ -262,7 +262,7 @@ func TestFilesystem(t *testing.T) {
 }
 
 func chtimesTree(tb testing.TB, dir string, atime, mtime time.Time) {
- is.New(tb).NoErr(filepath.WalkDir(dir, func(path string, d fs.DirEntry, err error) error {
+ is.New(tb).NoErr(filepath.WalkDir(dir, func(path string, _ fs.DirEntry, err error) error {
  if err != nil {
  return err
  }