github.com/cloud-foundations/dominator@v0.0.0-20221004181915-6e4fee580046/lib/srpc/proxy/impl.go (about)

     1  package proxy
     2  
     3  import (
     4  	"io"
     5  	"net"
     6  
     7  	"github.com/Cloud-Foundations/Dominator/lib/errors"
     8  	"github.com/Cloud-Foundations/Dominator/lib/srpc"
     9  	proto "github.com/Cloud-Foundations/Dominator/proto/proxy"
    10  )
    11  
    12  type autoFlusher struct {
    13  	writer writeFlusher
    14  }
    15  
    16  type writeFlusher interface {
    17  	Flush() error
    18  	io.Writer
    19  }
    20  
    21  func (t *srpcType) Connect(conn *srpc.Conn) error {
    22  	var request proto.ConnectRequest
    23  	if err := conn.Decode(&request); err != nil {
    24  		return err
    25  	}
    26  	requestedConn, err := net.DialTimeout(request.Network, request.Address,
    27  		request.Timeout)
    28  	e := conn.Encode(proto.ConnectResponse{Error: errors.ErrorToString(err)})
    29  	if e != nil {
    30  		return e
    31  	}
    32  	defer requestedConn.Close()
    33  	if err != nil {
    34  		return nil
    35  	}
    36  	if err := conn.Flush(); err != nil {
    37  		return err
    38  	}
    39  	t.logger.Debugf(0, "made proxy connection to: %s\n", request.Address)
    40  	closed := false
    41  	go func() {
    42  		defer requestedConn.Close()
    43  		if _, err := io.Copy(requestedConn, conn); err != nil {
    44  			if !closed {
    45  				t.logger.Printf("error copying proxied data to %s: %s\n",
    46  					request.Address, err)
    47  			}
    48  		}
    49  		closed = true
    50  	}()
    51  	if _, err := io.Copy(&autoFlusher{conn}, requestedConn); err != nil {
    52  		if !closed {
    53  			t.logger.Printf("error copying proxied data from %s: %s\n",
    54  				request.Address, err)
    55  		}
    56  	}
    57  	closed = true
    58  	return srpc.ErrorCloseClient
    59  }
    60  
    61  func (w *autoFlusher) Write(b []byte) (int, error) {
    62  	if nWritten, err := w.writer.Write(b); err != nil {
    63  		w.writer.Flush()
    64  		return nWritten, err
    65  	} else {
    66  		return nWritten, w.writer.Flush()
    67  	}
    68  }