github.com/mutagen-io/mutagen@v0.18.0-rc1/pkg/forwarding/forwarding.go (about)

     1  package forwarding
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  
     8  	"github.com/mutagen-io/mutagen/pkg/stream"
     9  )
    10  
    11  // ForwardAndClose performs bidirectional forwarding between the specified
    12  // connections. It waits for both directions to see EOF, for one direction to
    13  // see an error, or for context cancellation. Once one of these events occurs,
    14  // the connections are closed (terminating forwarding) and the function returns.
    15  // Both connections must implement CloseWriter or this function will panic. If
    16  // the caller passes non-nil values for firstAuditor and/or secondAuditor, then
    17  // auditing will be performed on the write end of the respective connection.
    18  func ForwardAndClose(ctx context.Context, first, second net.Conn, firstAuditor, secondAuditor stream.Auditor) {
    19  	// Defer closure of the connections.
    20  	defer func() {
    21  		first.Close()
    22  		second.Close()
    23  	}()
    24  
    25  	// Extract write closure interfaces.
    26  	firstCloseWriter, ok := first.(stream.CloseWriter)
    27  	if !ok {
    28  		panic("first connection does not implement write closure")
    29  	}
    30  	secondCloseWriter, ok := second.(stream.CloseWriter)
    31  	if !ok {
    32  		panic("second connection does not implement write closure")
    33  	}
    34  
    35  	// Forward traffic between the connections (with optional auditing) in
    36  	// separate Goroutines and track their termination. We track their
    37  	// termination via the error result, though this may be nil in the event
    38  	// that the source indicates EOF. If we do see an EOF from a source, then
    39  	// perform write closure on the corresponding destination in order to
    40  	// forward the EOF.
    41  	copyErrors := make(chan error, 2)
    42  	go func() {
    43  		_, err := io.Copy(stream.NewAuditWriter(first, firstAuditor), second)
    44  		if err == nil {
    45  			firstCloseWriter.CloseWrite()
    46  		}
    47  		copyErrors <- err
    48  	}()
    49  	go func() {
    50  		_, err := io.Copy(stream.NewAuditWriter(second, secondAuditor), first)
    51  		if err == nil {
    52  			secondCloseWriter.CloseWrite()
    53  		}
    54  		copyErrors <- err
    55  	}()
    56  
    57  	// Wait for both forwarding routines to finish while also monitoring for
    58  	// termination. We only abort this wait if we see a non-nil copy error from
    59  	// one of the forwarding routines (or forwarding is terminated). We allow
    60  	// nil errors because they simply indicate EOF and can be sent by some
    61  	// connection types by performing a half-close of a stream.
    62  	for i := 0; i < 2; i++ {
    63  		select {
    64  		case err := <-copyErrors:
    65  			if err != nil {
    66  				return
    67  			}
    68  		case <-ctx.Done():
    69  			return
    70  		}
    71  	}
    72  }