tractor.dev/toolkit-go@v0.0.0-20241010005851-214d91207d07/duplex/mux/proxy.go (about)

     1  package mux
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"sync"
     7  )
     8  
     9  // Proxy accepts channels on src then opens a channel on dst and performs
    10  // an io.Copy in both directions in goroutines. Proxy returns non-EOF errors
    11  // from src.Accept, nil on EOF, and any errors from dst.Open after closing
    12  // the accepted channel from src.
    13  func Proxy(dst, src Session) error {
    14  	for {
    15  		ctx := context.Background()
    16  		a, err := src.Accept()
    17  		if err != nil {
    18  			if err == io.EOF {
    19  				return nil
    20  			}
    21  			return err
    22  		}
    23  		b, err := dst.Open(ctx)
    24  		if err != nil {
    25  			a.Close()
    26  			return err
    27  		}
    28  		go proxy(a, b)
    29  	}
    30  }
    31  
    32  func proxy(a, b Channel) {
    33  	var wg sync.WaitGroup
    34  	wg.Add(2)
    35  	go func() {
    36  		io.Copy(a, b)
    37  		a.CloseWrite()
    38  		wg.Done()
    39  	}()
    40  	go func() {
    41  		io.Copy(b, a)
    42  		b.CloseWrite()
    43  		wg.Done()
    44  	}()
    45  	wg.Wait()
    46  	a.Close()
    47  	b.Close()
    48  }