Added allowed IPs support
This commit is contained in:
+32
@@ -55,6 +55,7 @@ type Config struct {
|
||||
type Server struct {
|
||||
config *Config
|
||||
authMethods map[uint8]Authenticator
|
||||
isIPAllowed func(net.IP) bool
|
||||
}
|
||||
|
||||
// New creates a new Server and potentially returns an error
|
||||
@@ -93,6 +94,11 @@ func New(conf *Config) (*Server, error) {
|
||||
server.authMethods[a.GetCode()] = a
|
||||
}
|
||||
|
||||
// Set default IP whitelist function
|
||||
server.isIPAllowed = func(ip net.IP) bool {
|
||||
return true // default allow all IPs
|
||||
}
|
||||
|
||||
return server, nil
|
||||
}
|
||||
|
||||
@@ -117,11 +123,37 @@ func (s *Server) Serve(l net.Listener) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// SetIPWhitelist sets the function to check if a given IP is allowed
|
||||
func (s *Server) SetIPWhitelist(allowedIPs []net.IP) {
|
||||
s.isIPAllowed = func(ip net.IP) bool {
|
||||
for _, allowedIP := range allowedIPs {
|
||||
if ip.Equal(allowedIP) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// ServeConn is used to serve a single connection.
|
||||
func (s *Server) ServeConn(conn net.Conn) error {
|
||||
defer conn.Close()
|
||||
bufConn := bufio.NewReader(conn)
|
||||
|
||||
// Check client IP against whitelist
|
||||
clientIP, _, err := net.SplitHostPort(conn.RemoteAddr().String())
|
||||
if err != nil {
|
||||
s.config.Logger.Printf("[ERR] socks: Failed to get client IP address: %v", err)
|
||||
return err
|
||||
}
|
||||
ip := net.ParseIP(clientIP)
|
||||
if s.isIPAllowed(ip) {
|
||||
s.config.Logger.Printf("[INFO] socks: Connection from allowed IP address: %s", clientIP)
|
||||
} else {
|
||||
s.config.Logger.Printf("[WARN] socks: Connection from not allowed IP address: %s", clientIP)
|
||||
return fmt.Errorf("connection from not allowed IP address")
|
||||
}
|
||||
|
||||
// Read the version byte
|
||||
version := []byte{0}
|
||||
if _, err := bufConn.Read(version); err != nil {
|
||||
|
||||
Reference in New Issue
Block a user