134 lines
2.8 KiB
Go
134 lines
2.8 KiB
Go
package main
|
||
|
||
import (
|
||
"archive/tar"
|
||
"flag"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"os"
|
||
"path/filepath"
|
||
"strconv"
|
||
"strings"
|
||
|
||
"github.com/google/renameio/v2"
|
||
"golang.org/x/crypto/ssh"
|
||
)
|
||
|
||
type countingWriter int64
|
||
|
||
func (cw *countingWriter) Write(p []byte) (n int, err error) {
|
||
*cw += countingWriter(len(p))
|
||
return len(p), nil
|
||
}
|
||
|
||
func scpSink(channel ssh.Channel, req *ssh.Request, cmdline []string) error {
|
||
scpFlags := flag.NewFlagSet("scp", flag.ContinueOnError)
|
||
sink := scpFlags.Bool("t", false, "sink (to)")
|
||
if err := scpFlags.Parse(cmdline[1:]); err != nil {
|
||
return err
|
||
}
|
||
if !*sink {
|
||
return fmt.Errorf("expected -t")
|
||
}
|
||
|
||
// Tell the remote end we’re ready to receive data.
|
||
if _, err := channel.Write([]byte{0x00}); err != nil {
|
||
return err
|
||
}
|
||
|
||
buf := make([]byte, 1024)
|
||
for {
|
||
n, err := channel.Read(buf)
|
||
if err == io.EOF {
|
||
break
|
||
}
|
||
if err != nil {
|
||
return err
|
||
}
|
||
msg := buf[:n]
|
||
|
||
// Acknowledge receipt of the control message
|
||
if _, err := channel.Write([]byte{0x00}); err != nil {
|
||
return err
|
||
}
|
||
|
||
if msg[0] == 'C' {
|
||
msgstr := strings.TrimSpace(string(msg))
|
||
parts := strings.Split(msgstr, " ")
|
||
if got, want := len(parts), 3; got != want {
|
||
return fmt.Errorf("invalid number of space-separated tokens in control message %q: got %d, want %d", msgstr, got, want)
|
||
}
|
||
size, err := strconv.ParseInt(parts[1], 0, 64)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
// Retrieve file contents
|
||
var cw countingWriter
|
||
if err := unpackTar(io.TeeReader(channel, &cw)); err != nil {
|
||
return err
|
||
}
|
||
if rest := size - int64(cw); rest > 0 {
|
||
buf := make([]byte, rest)
|
||
if _, err := channel.Read(buf); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
|
||
// Read status byte after transfer
|
||
buf := make([]byte, 1)
|
||
if _, err := channel.Read(buf); err != nil {
|
||
return err
|
||
}
|
||
|
||
// Acknowledge file transfer
|
||
if _, err := channel.Write([]byte{0x00}); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
}
|
||
|
||
exitStatus := make([]byte, 4)
|
||
exitStatus[3] = 0
|
||
if _, err := channel.SendRequest("exit-status", false, exitStatus); err != nil {
|
||
return err
|
||
}
|
||
channel.Close()
|
||
req.Reply(true, nil)
|
||
return nil
|
||
}
|
||
|
||
func unpackTar(r io.Reader) error {
|
||
tr := tar.NewReader(r)
|
||
for {
|
||
h, err := tr.Next()
|
||
if err == io.EOF {
|
||
break
|
||
}
|
||
if err != nil {
|
||
return err
|
||
}
|
||
|
||
log.Printf("extracting %q", h.Name)
|
||
if err := os.MkdirAll(filepath.Dir(h.Name), 0700); err != nil {
|
||
return err
|
||
}
|
||
if strings.HasSuffix(h.Name, "/") {
|
||
continue // directory, don’t try to OpenFile() it
|
||
}
|
||
mode := h.FileInfo().Mode() & os.ModePerm
|
||
out, err := renameio.NewPendingFile(h.Name, renameio.WithStaticPermissions(mode))
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if _, err := io.Copy(out, tr); err != nil {
|
||
return err
|
||
}
|
||
if err := out.CloseAtomicallyReplace(); err != nil {
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|