Пример #1
0
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
}
Пример #2
0
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...)
}
Пример #3
0
// 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
}
Пример #4
0
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())
	}
}
Пример #5
0
// 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
}
Пример #6
0
// 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
}