func (w *Worker) doWrite() error {
filename := w.getFilename()
file, err := w.fs.Create(filename)
if err != nil {
return err
}
flow := flowrate.NewWriter(file, -1)
var size int64
expected := w.getSize()
for {
s, err := copyN(flow, w.r, w.c.BlockSize)
if err != nil {
return err
}
size += int64(s)
if size >= expected {
break
}
}
w.filenames = append(w.filenames, filename)
flow.Close()
w.WStatus.Add(NewStatus(flow.Status()))
return nil
}
func NewClient(conn *ssh.Client, readLimitBytesPerSecond, writeLimitBytesPerSecond int64, opts ...func(*sftp.Client) error) (*sftp.Client, error) {
s, err := conn.NewSession()
if err != nil {
return nil, err
}
if err := s.RequestSubsystem("sftp"); err != nil {
return nil, err
}
pw, err := s.StdinPipe()
if err != nil {
return nil, err
}
pr, err := s.StdoutPipe()
if err != nil {
return nil, err
}
if readLimitBytesPerSecond > 0 {
pr = flowrate.NewReader(pr, readLimitBytesPerSecond)
}
if writeLimitBytesPerSecond > 0 {
pw = flowrate.NewWriter(pw, writeLimitBytesPerSecond)
}
return sftp.NewClientPipe(pr, pw, opts...)
}
// NewDownloader creates new instance of Downloader which specified number
// of threads and download limit in bytes/sec
func NewDownloader(threads int, downLimit int64, progress aptly.Progress) aptly.Downloader {
transport := *http.DefaultTransport.(*http.Transport)
transport.DisableCompression = true
transport.RegisterProtocol("ftp", &protocol.FTPRoundTripper{})
downloader := &downloaderImpl{
queue: make(chan *downloadTask, 1000),
stop: make(chan struct{}, threads),
stopped: make(chan struct{}, threads),
pause: make(chan struct{}),
unpause: make(chan struct{}),
threads: threads,
progress: progress,
client: &http.Client{
Transport: &transport,
},
}
if downLimit > 0 {
downloader.aggWriter = flowrate.NewWriter(progress, downLimit)
} else {
downloader.aggWriter = progress
}
for i := 0; i < downloader.threads; i++ {
go downloader.process()
}
return downloader
}
func sendFileToRemoteHost(client *ssh.Client, limit int64, sourceFile, targetUser, targetHost, targetFile string) {
session, err := client.NewSession()
if err != nil {
log.Fatalln("Failed to create session: " + err.Error())
}
defer session.Close()
go func() {
iw, err := session.StdinPipe()
if err != nil {
log.Fatalln("Failed to create input pipe: " + err.Error())
}
w := flowrate.NewWriter(iw, limit)
src, srcErr := os.Open(sourceFile)
if srcErr != nil {
log.Fatalln("Failed to open source file: " + srcErr.Error())
}
srcStat, statErr := src.Stat()
if statErr != nil {
log.Fatalln("Failed to stat file: " + statErr.Error())
}
fmt.Fprintln(w, "C0644", srcStat.Size(), filepath.Base(sourceFile))
if srcStat.Size() > 0 {
bar := pb.New(int(srcStat.Size()))
bar.Units = pb.U_BYTES
bar.ShowSpeed = true
bar.Start()
wp := io.MultiWriter(w, bar)
fmt.Printf("Transferring %s to %s@%s:%s\n", sourceFile, targetUser, targetHost, targetFile)
fmt.Printf("Speed limited to %d bytes/sec\n", limit)
io.Copy(wp, src)
bar.Finish()
fmt.Fprint(w, "\x00")
w.Close()
} else {
fmt.Printf("Transferred empty file %s to %s@%s:%s\n", sourceFile, targetUser, targetHost, targetFile)
fmt.Fprint(w, "\x00")
w.Close()
}
}()
if err := session.Run(fmt.Sprintf("scp -t %s", targetFile)); err != nil {
log.Fatalln("Failed to run: " + err.Error())
}
}
// tryUpgrade returns true if the request was handled.
func (h *UpgradeAwareProxyHandler) tryUpgrade(w http.ResponseWriter, req *http.Request) bool {
if !httpstream.IsUpgradeRequest(req) {
return false
}
backendConn, err := proxy.DialURL(h.Location, h.Transport)
if err != nil {
h.Responder.Error(err)
return true
}
defer backendConn.Close()
requestHijackedConn, _, err := w.(http.Hijacker).Hijack()
if err != nil {
h.Responder.Error(err)
return true
}
defer requestHijackedConn.Close()
newReq, err := http.NewRequest(req.Method, h.Location.String(), req.Body)
if err != nil {
h.Responder.Error(err)
return true
}
newReq.Header = req.Header
if err = newReq.Write(backendConn); err != nil {
h.Responder.Error(err)
return true
}
wg := &sync.WaitGroup{}
wg.Add(2)
go func() {
var writer io.WriteCloser
if h.MaxBytesPerSec > 0 {
writer = flowrate.NewWriter(backendConn, h.MaxBytesPerSec)
} else {
writer = backendConn
}
_, err := io.Copy(writer, requestHijackedConn)
if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
glog.Errorf("Error proxying data from client to backend: %v", err)
}
wg.Done()
}()
go func() {
var reader io.ReadCloser
if h.MaxBytesPerSec > 0 {
reader = flowrate.NewReader(backendConn, h.MaxBytesPerSec)
} else {
reader = backendConn
}
_, err := io.Copy(requestHijackedConn, reader)
if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
glog.Errorf("Error proxying data from backend to client: %v", err)
}
wg.Done()
}()
wg.Wait()
return true
}
// tryUpgrade returns true if the request was handled.
func (h *UpgradeAwareProxyHandler) tryUpgrade(w http.ResponseWriter, req *http.Request) bool {
if !httpstream.IsUpgradeRequest(req) {
return false
}
var (
backendConn net.Conn
rawResponse []byte
err error
)
if h.InterceptRedirects && utilconfig.DefaultFeatureGate.StreamingProxyRedirects() {
backendConn, rawResponse, err = h.connectBackendWithRedirects(req)
} else {
backendConn, err = h.connectBackend(req.Method, h.Location, req.Header, req.Body)
}
if err != nil {
h.Responder.Error(err)
return true
}
defer backendConn.Close()
// Once the connection is hijacked, the ErrorResponder will no longer work, so
// hijacking should be the last step in the upgrade.
requestHijacker, ok := w.(http.Hijacker)
if !ok {
h.Responder.Error(fmt.Errorf("request connection cannot be hijacked: %T", w))
return true
}
requestHijackedConn, _, err := requestHijacker.Hijack()
if err != nil {
h.Responder.Error(fmt.Errorf("error hijacking request connection: %v", err))
return true
}
defer requestHijackedConn.Close()
// Forward raw response bytes back to client.
if len(rawResponse) > 0 {
if _, err = requestHijackedConn.Write(rawResponse); err != nil {
utilruntime.HandleError(fmt.Errorf("Error proxying response from backend to client: %v", err))
}
}
// Proxy the connection.
wg := &sync.WaitGroup{}
wg.Add(2)
go func() {
var writer io.WriteCloser
if h.MaxBytesPerSec > 0 {
writer = flowrate.NewWriter(backendConn, h.MaxBytesPerSec)
} else {
writer = backendConn
}
_, err := io.Copy(writer, requestHijackedConn)
if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
glog.Errorf("Error proxying data from client to backend: %v", err)
}
wg.Done()
}()
go func() {
var reader io.ReadCloser
if h.MaxBytesPerSec > 0 {
reader = flowrate.NewReader(backendConn, h.MaxBytesPerSec)
} else {
reader = backendConn
}
_, err := io.Copy(requestHijackedConn, reader)
if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
glog.Errorf("Error proxying data from backend to client: %v", err)
}
wg.Done()
}()
wg.Wait()
return true
}
