// Copyright 2020, Chef. All rights reserved. // https://github.com/q191201771/naza // // Use of this source code is governed by a MIT-style license // that can be found in the License file. // // Author: Chef (191201771@qq.com) package nazanet import ( "net" "sync" ) // 从指定的UDP端口范围内,寻找可绑定监听的端口,绑定监听并返回 // Pool只提供Acquire获取接口,不提供释放接口,连接资源是标准*net.UDPConn对象,需要释放时,外部直接Close即可 // type AvailUDPConnPool struct { minPort uint16 maxPort uint16 m sync.Mutex lastPort uint16 } func NewAvailUDPConnPool(minPort uint16, maxPort uint16) *AvailUDPConnPool { return &AvailUDPConnPool{ minPort: minPort, maxPort: maxPort, lastPort: minPort, } } func (a *AvailUDPConnPool) Acquire() (*net.UDPConn, uint16, error) { a.m.Lock() defer a.m.Unlock() loopFirstFlag := true p := a.lastPort for { // 找了一轮也没有可用的,返回错误 if !loopFirstFlag && p == a.lastPort { return nil, 0, ErrNazaNet } loopFirstFlag = false conn, err := listenUDPWithPort(p) // 绑定失败,尝试下一个端口 if err != nil { p = a.nextPort(p) continue } // 绑定成功,更新last,返回结果 a.lastPort = a.nextPort(p) return conn, p, nil } } // 有的业务场景,需要返回两个可用的端口,并且必须是连续的 // @return 前面的是端口小的,后面的是端口+1的 // func (a *AvailUDPConnPool) Acquire2() (*net.UDPConn, uint16, *net.UDPConn, uint16, error) { a.m.Lock() defer a.m.Unlock() loopFirstFlag := true p := a.lastPort for { // 找了一轮也没有可用的,返回错误 if !loopFirstFlag && p == a.lastPort { return nil, 0, nil, 0, ErrNazaNet } loopFirstFlag = false // 因为第一个端口如果为最大值,那么和第二个端口肯定不是线性连续了 if p == a.maxPort { p = a.minPort continue } conn, err := listenUDPWithPort(p) // 第一个就绑定失败,尝试下一个端口 if err != nil { p = a.nextPort(p) continue } // 绑定成功,因为我们需要两个,所以我们还要找第二个 // 因为前面已经有判断最大值了,所以直接+1 conn2, err := listenUDPWithPort(p + 1) // 第二个失败了,关闭第一个,然后从第二个的下一个重新尝试 if err != nil { _ = conn.Close() p = a.nextPort(p + 1) continue } // 绑定成功,更新last,返回结果 a.lastPort = a.nextPort(p + 1) return conn, p, conn2, p + 1, nil } } // 通过Acquire获取到可用net.UDPConn对象后,将对象关闭,只返回可用的端口 func (a *AvailUDPConnPool) Peek() (uint16, error) { conn, port, err := a.Acquire() if err == nil { err = conn.Close() } return port, err } func (a *AvailUDPConnPool) nextPort(p uint16) uint16 { if p == a.maxPort { return a.minPort } return p + 1 }