github.com/mysteriumnetwork/node@v0.0.0-20240516044423-365054f76801/router/protect.go (about)

     1  /*
     2   * Copyright (C) 2021 The "MysteriumNetwork/node" Authors.
     3   *
     4   * This program is free software: you can redistribute it and/or modify
     5   * it under the terms of the GNU General Public License as published by
     6   * the Free Software Foundation, either version 3 of the License, or
     7   * (at your option) any later version.
     8   *
     9   * This program is distributed in the hope that it will be useful,
    10   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    11   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    12   * GNU General Public License for more details.
    13   *
    14   * You should have received a copy of the GNU General Public License
    15   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    16   */
    17  
    18  package router
    19  
    20  import (
    21  	"net"
    22  	"sync"
    23  )
    24  
    25  var (
    26  	mu      sync.RWMutex
    27  	protect func(fd int) error = nil
    28  )
    29  
    30  // SetProtectFunc sets the callback for using to protect provided connection from going through the tunnel.
    31  func SetProtectFunc(f func(fd int) error) {
    32  	mu.Lock()
    33  	defer mu.Unlock()
    34  
    35  	protect = f
    36  }
    37  
    38  // Protect protects provided connection from going through the tunnel.
    39  func Protect(fd int) error {
    40  	mu.RLock()
    41  	defer mu.RUnlock()
    42  
    43  	if protect == nil {
    44  		return nil
    45  	}
    46  
    47  	return protect(fd)
    48  }
    49  
    50  // ProtectUDPConn protects provided UDP connection from going through the tunnel.
    51  func ProtectUDPConn(c *net.UDPConn) error {
    52  	sysconn, err := c.SyscallConn()
    53  	if err != nil {
    54  		return err
    55  	}
    56  
    57  	fd := 0
    58  
    59  	err = sysconn.Control(func(f uintptr) {
    60  		fd = int(f)
    61  	})
    62  	if err != nil {
    63  		return err
    64  	}
    65  
    66  	return Protect(fd)
    67  }