github.com/Bytom/bytom@v1.1.2-0.20210127130405-ae40204c0b09/p2p/upnp/probe.go (about)

     1  package upnp
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"net"
     7  	"time"
     8  
     9  	log "github.com/sirupsen/logrus"
    10  )
    11  
    12  type UPNPCapabilities struct {
    13  	PortMapping bool
    14  	Hairpin     bool
    15  }
    16  
    17  func makeUPNPListener(intPort int, extPort int) (NAT, net.Listener, net.IP, error) {
    18  	nat, err := Discover()
    19  	if err != nil {
    20  		return nil, nil, nil, errors.New(fmt.Sprintf("NAT upnp could not be discovered: %v", err))
    21  	}
    22  	log.WithField("ourIP", nat.(*upnpNAT).ourIP).Info("outIP:")
    23  
    24  	ext, err := nat.GetExternalAddress()
    25  	if err != nil {
    26  		return nat, nil, nil, errors.New(fmt.Sprintf("External address error: %v", err))
    27  	}
    28  	log.WithField("address", ext).Info("External address")
    29  
    30  	port, err := nat.AddPortMapping("tcp", extPort, intPort, "Tendermint UPnP Probe", 0)
    31  	if err != nil {
    32  		return nat, nil, ext, errors.New(fmt.Sprintf("Port mapping error: %v", err))
    33  	}
    34  	log.WithField("port", port).Info("Port mapping mapped")
    35  
    36  	// also run the listener, open for all remote addresses.
    37  	listener, err := net.Listen("tcp", fmt.Sprintf(":%v", intPort))
    38  	if err != nil {
    39  		return nat, nil, ext, errors.New(fmt.Sprintf("Error establishing listener: %v", err))
    40  	}
    41  	return nat, listener, ext, nil
    42  }
    43  
    44  func testHairpin(listener net.Listener, extAddr string) (supportsHairpin bool) {
    45  	// Listener
    46  	go func() {
    47  		inConn, err := listener.Accept()
    48  		if err != nil {
    49  			log.WithField("error", err).Error("Listener.Accept() error")
    50  			return
    51  		}
    52  		log.WithFields(log.Fields{
    53  			"LocalAddr":  inConn.LocalAddr(),
    54  			"RemoteAddr": inConn.RemoteAddr(),
    55  		}).Info("Accepted incoming connection")
    56  		buf := make([]byte, 1024)
    57  		n, err := inConn.Read(buf)
    58  		if err != nil {
    59  			log.WithField("error", err).Error("Incoming connection read error")
    60  			return
    61  		}
    62  		log.Infof("Incoming connection read %v bytes: %X", n, buf)
    63  		if string(buf) == "test data" {
    64  			supportsHairpin = true
    65  			return
    66  		}
    67  	}()
    68  
    69  	// Establish outgoing
    70  	outConn, err := net.Dial("tcp", extAddr)
    71  	if err != nil {
    72  		log.WithField("error", err).Error("Outgoing connection dial error")
    73  		return
    74  	}
    75  
    76  	n, err := outConn.Write([]byte("test data"))
    77  	if err != nil {
    78  		log.WithField("error", err).Error("Outgoing connection write error")
    79  		return
    80  	}
    81  	log.Infof("Outgoing connection wrote %v bytes", n)
    82  
    83  	// Wait for data receipt
    84  	time.Sleep(1 * time.Second)
    85  	return
    86  }
    87  
    88  func Probe() (caps UPNPCapabilities, err error) {
    89  	log.Info("Probing for UPnP!")
    90  
    91  	intPort, extPort := 8001, 8001
    92  
    93  	nat, listener, ext, err := makeUPNPListener(intPort, extPort)
    94  	if err != nil {
    95  		return
    96  	}
    97  	caps.PortMapping = true
    98  
    99  	// Deferred cleanup
   100  	defer func() {
   101  		err = nat.DeletePortMapping("tcp", intPort, extPort)
   102  		if err != nil {
   103  			log.WithField("error", err).Error("Port mapping delete error")
   104  		}
   105  		listener.Close()
   106  	}()
   107  
   108  	supportsHairpin := testHairpin(listener, fmt.Sprintf("%v:%v", ext, extPort))
   109  	if supportsHairpin {
   110  		caps.Hairpin = true
   111  	}
   112  
   113  	return
   114  }