github.com/lastbackend/toolkit@v0.0.0-20241020043710-cafa37b95aad/pkg/util/net/net.go (about)

     1  /*
     2  Copyright [2014] - [2023] The Last.Backend authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package net
    18  
    19  import (
    20  	"errors"
    21  	"fmt"
    22  	"net"
    23  	"strconv"
    24  	"strings"
    25  )
    26  
    27  // HostPort format addr and port suitable for dial
    28  func HostPort(addr string, port interface{}) string {
    29  	host := addr
    30  	if strings.Count(addr, ":") > 0 {
    31  		host = fmt.Sprintf("[%s]", addr)
    32  	}
    33  	// when port is blank or 0, host is a queue name
    34  	if v, ok := port.(string); ok && v == "" {
    35  		return host
    36  	} else if v, ok := port.(int); ok && v == 0 && net.ParseIP(host) == nil {
    37  		return host
    38  	}
    39  
    40  	return fmt.Sprintf("%s:%v", host, port)
    41  }
    42  
    43  // Listen takes addr:portmin-portmax and binds to the first available port
    44  // Example: Listen("localhost:5000-6000", fn)
    45  func Listen(addr string, fn func(string) (net.Listener, error)) (net.Listener, error) {
    46  
    47  	if strings.Count(addr, ":") == 1 && strings.Count(addr, "-") == 0 {
    48  		return fn(addr)
    49  	}
    50  
    51  	// host:port || host:min-max
    52  	host, ports, err := net.SplitHostPort(addr)
    53  	if err != nil {
    54  		return nil, err
    55  	}
    56  
    57  	// try to extract port range
    58  	prange := strings.Split(ports, "-")
    59  
    60  	// single port
    61  	if len(prange) < 2 {
    62  		return fn(addr)
    63  	}
    64  
    65  	// we have a port range
    66  
    67  	// extract min port
    68  	min, err := strconv.Atoi(prange[0])
    69  	if err != nil {
    70  		return nil, errors.New("unable to extract port range")
    71  	}
    72  
    73  	// extract max port
    74  	max, err := strconv.Atoi(prange[1])
    75  	if err != nil {
    76  		return nil, errors.New("unable to extract port range")
    77  	}
    78  
    79  	// range the ports
    80  	for port := min; port <= max; port++ {
    81  		// try bind to host:port
    82  		ln, err := fn(HostPort(host, port))
    83  		if err == nil {
    84  			return ln, nil
    85  		}
    86  
    87  		// hit max port
    88  		if port == max {
    89  			return nil, err
    90  		}
    91  	}
    92  
    93  	// why are we here?
    94  	return nil, fmt.Errorf("unable to bind to %s", addr)
    95  }