github.com/bepass-org/wireguard-go@v1.0.4-rc2.0.20240304192354-ebce6572bc24/wiresocks/proxy.go (about)

     1  package wiresocks
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"log"
     7  	"net/netip"
     8  	"time"
     9  
    10  	"github.com/bepass-org/proxy/pkg/mixed"
    11  	"github.com/bepass-org/proxy/pkg/statute"
    12  	"github.com/bepass-org/wireguard-go/device"
    13  	"github.com/bepass-org/wireguard-go/tun/netstack"
    14  )
    15  
    16  // VirtualTun stores a reference to netstack network and DNS configuration
    17  type VirtualTun struct {
    18  	Tnet      *netstack.Net
    19  	SystemDNS bool
    20  	Verbose   bool
    21  	Logger    DefaultLogger
    22  	Dev       *device.Device
    23  	Ctx       context.Context
    24  }
    25  
    26  type DefaultLogger struct {
    27  	verbose bool
    28  }
    29  
    30  func (l DefaultLogger) Debug(v ...interface{}) {
    31  	if l.verbose {
    32  		log.Println(v...)
    33  	}
    34  }
    35  
    36  func (l DefaultLogger) Error(v ...interface{}) {
    37  	log.Println(v...)
    38  }
    39  
    40  // StartProxy spawns a socks5 server.
    41  func (vt *VirtualTun) StartProxy(bindAddress netip.AddrPort) {
    42  	proxy := mixed.NewProxy(
    43  		mixed.WithBindAddress(bindAddress.String()),
    44  		mixed.WithLogger(vt.Logger),
    45  		mixed.WithContext(vt.Ctx),
    46  		mixed.WithUserHandler(func(request *statute.ProxyRequest) error {
    47  			return vt.generalHandler(request)
    48  		}),
    49  	)
    50  	go func() {
    51  		_ = proxy.ListenAndServe()
    52  	}()
    53  	go func() {
    54  		for {
    55  			select {
    56  			case <-vt.Ctx.Done():
    57  				vt.Stop()
    58  				return
    59  			default:
    60  				time.Sleep(500 * time.Millisecond)
    61  			}
    62  		}
    63  	}()
    64  }
    65  
    66  func (vt *VirtualTun) generalHandler(req *statute.ProxyRequest) error {
    67  	if vt.Verbose {
    68  		log.Printf("handling %s request to %s", req.Network, req.Destination)
    69  	}
    70  	conn, err := vt.Tnet.Dial(req.Network, req.Destination)
    71  	if err != nil {
    72  		return err
    73  	}
    74  	// Close the connections when this function exits
    75  	defer conn.Close()
    76  	defer req.Conn.Close()
    77  	// Channel to notify when copy operation is done
    78  	done := make(chan error, 1)
    79  	// Copy data from req.Conn to conn
    80  	go func() {
    81  		_, err := io.Copy(conn, req.Conn)
    82  		done <- err
    83  	}()
    84  	// Copy data from conn to req.Conn
    85  	go func() {
    86  		_, err := io.Copy(req.Conn, conn)
    87  		done <- err
    88  	}()
    89  	// Wait for one of the copy operations to finish
    90  	err = <-done
    91  	if err != nil {
    92  		log.Println(err)
    93  	}
    94  	// Close connections and wait for the other copy operation to finish
    95  	conn.Close()
    96  	req.Conn.Close()
    97  	<-done
    98  
    99  	return nil
   100  }
   101  
   102  func (vt *VirtualTun) Stop() {
   103  	if vt.Dev != nil {
   104  		err := vt.Dev.Down()
   105  		if err != nil {
   106  			log.Println(err)
   107  		}
   108  	}
   109  }