115 lines
2.5 KiB
Go
115 lines
2.5 KiB
Go
package portforward
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/huin/goupnp/dcps/internetgateway1"
|
|
"github.com/veggiedefender/torrent-client/internal/logger"
|
|
)
|
|
|
|
type PortManager struct {
|
|
mu sync.Mutex
|
|
clients []*internetgateway1.WANIPConnection1
|
|
port uint16
|
|
opened bool
|
|
}
|
|
|
|
func NewPortManager() *PortManager {
|
|
return &PortManager{}
|
|
}
|
|
|
|
// OpenPort attempts to open the specified port on all discovered UPnP routers.
|
|
func (pm *PortManager) OpenPort(ctx context.Context, port uint16, desc string) error {
|
|
pm.mu.Lock()
|
|
defer pm.mu.Unlock()
|
|
|
|
if pm.opened && pm.port == port {
|
|
return nil
|
|
}
|
|
|
|
logger.Info("SYS", "Discovering UPnP routers for port %d...", port)
|
|
|
|
discoverCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
|
|
defer cancel()
|
|
|
|
discoverDone := make(chan []*internetgateway1.WANIPConnection1, 1)
|
|
go func() {
|
|
clients, _, _ := internetgateway1.NewWANIPConnection1Clients()
|
|
discoverDone <- clients
|
|
}()
|
|
|
|
var clients []*internetgateway1.WANIPConnection1
|
|
select {
|
|
case <-discoverCtx.Done():
|
|
logger.Warn("SYS", "UPnP discovery timed out")
|
|
return nil
|
|
case clients = <-discoverDone:
|
|
}
|
|
|
|
if len(clients) == 0 {
|
|
logger.Warn("SYS", "No UPnP routers discovered")
|
|
return nil
|
|
}
|
|
|
|
localIP, err := getLocalIP()
|
|
if err != nil {
|
|
logger.Warn("SYS", "Could not get local IP for UPnP")
|
|
return nil
|
|
}
|
|
|
|
success := false
|
|
for _, client := range clients {
|
|
// TCP Mapping
|
|
errTCP := client.AddPortMapping("", port, "TCP", port, localIP, true, desc, 0)
|
|
// UDP Mapping (for DHT)
|
|
errUDP := client.AddPortMapping("", port, "UDP", port, localIP, true, desc, 0)
|
|
|
|
if errTCP == nil || errUDP == nil {
|
|
success = true
|
|
}
|
|
}
|
|
|
|
if success {
|
|
pm.clients = clients
|
|
pm.port = port
|
|
pm.opened = true
|
|
logger.Info("SYS", "UPnP successfully forwarded port %d", port)
|
|
} else {
|
|
logger.Warn("SYS", "Failed to add UPnP port mapping")
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// ClosePort closes the currently opened port.
|
|
func (pm *PortManager) ClosePort() error {
|
|
pm.mu.Lock()
|
|
defer pm.mu.Unlock()
|
|
|
|
if !pm.opened || len(pm.clients) == 0 {
|
|
return nil
|
|
}
|
|
|
|
for _, client := range pm.clients {
|
|
_ = client.DeletePortMapping("", pm.port, "TCP")
|
|
_ = client.DeletePortMapping("", pm.port, "UDP")
|
|
}
|
|
|
|
logger.Info("SYS", "UPnP successfully cleared port %d", pm.port)
|
|
pm.opened = false
|
|
return nil
|
|
}
|
|
|
|
// getLocalIP returns the preferred outbound IP of this machine
|
|
func getLocalIP() (string, error) {
|
|
conn, err := net.Dial("udp", "8.8.8.8:80")
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer conn.Close()
|
|
localAddr := conn.LocalAddr().(*net.UDPAddr)
|
|
return localAddr.IP.String(), nil
|
|
}
|