update group ci

server/group/group.go

@@ -0,0 +1,25 @@
+// Copyright 2018 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package group
+
+import (
+	"errors"
+)
+
+var (
+	ErrGroupAuthFailed    = errors.New("group auth failed")
+	ErrGroupParamsInvalid = errors.New("group params invalid")
+	ErrListenerClosed     = errors.New("group listener closed")
+)

server/group/tcp.go

@@ -0,0 +1,200 @@
+// Copyright 2018 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package group
+
+import (
+	"fmt"
+	"net"
+	"sync"
+
+	"github.com/fatedier/frp/server/ports"
+
+	gerr "github.com/fatedier/golib/errors"
+)
+
+type TcpGroupListener struct {
+	groupName string
+	group     *TcpGroup
+
+	addr    net.Addr
+	closeCh chan struct{}
+}
+
+func newTcpGroupListener(name string, group *TcpGroup, addr net.Addr) *TcpGroupListener {
+	return &TcpGroupListener{
+		groupName: name,
+		group:     group,
+		addr:      addr,
+		closeCh:   make(chan struct{}),
+	}
+}
+
+func (ln *TcpGroupListener) Accept() (c net.Conn, err error) {
+	var ok bool
+	select {
+	case <-ln.closeCh:
+		return nil, ErrListenerClosed
+	case c, ok = <-ln.group.Accept():
+		if !ok {
+			return nil, ErrListenerClosed
+		}
+		return c, nil
+	}
+}
+
+func (ln *TcpGroupListener) Addr() net.Addr {
+	return ln.addr
+}
+
+func (ln *TcpGroupListener) Close() (err error) {
+	close(ln.closeCh)
+	ln.group.CloseListener(ln)
+	return
+}
+
+type TcpGroup struct {
+	group    string
+	groupKey string
+	addr     string
+	port     int
+	realPort int
+
+	acceptCh chan net.Conn
+	index    uint64
+	tcpLn    net.Listener
+	lns      []*TcpGroupListener
+	ctl      *TcpGroupCtl
+	mu       sync.Mutex
+}
+
+func NewTcpGroup(ctl *TcpGroupCtl) *TcpGroup {
+	return &TcpGroup{
+		lns:      make([]*TcpGroupListener, 0),
+		ctl:      ctl,
+		acceptCh: make(chan net.Conn),
+	}
+}
+
+func (tg *TcpGroup) Listen(proxyName string, group string, groupKey string, addr string, port int) (ln *TcpGroupListener, realPort int, err error) {
+	tg.mu.Lock()
+	defer tg.mu.Unlock()
+	if len(tg.lns) == 0 {
+		realPort, err = tg.ctl.portManager.Acquire(proxyName, port)
+		if err != nil {
+			return
+		}
+		tcpLn, errRet := net.Listen("tcp", fmt.Sprintf("%s:%d", addr, port))
+		if errRet != nil {
+			err = errRet
+			return
+		}
+		ln = newTcpGroupListener(group, tg, tcpLn.Addr())
+
+		tg.group = group
+		tg.groupKey = groupKey
+		tg.addr = addr
+		tg.port = port
+		tg.realPort = realPort
+		tg.tcpLn = tcpLn
+		tg.lns = append(tg.lns, ln)
+		if tg.acceptCh == nil {
+			tg.acceptCh = make(chan net.Conn)
+		}
+		go tg.worker()
+	} else {
+		if tg.group != group || tg.addr != addr || tg.port != port {
+			err = ErrGroupParamsInvalid
+			return
+		}
+		if tg.groupKey != groupKey {
+			err = ErrGroupAuthFailed
+			return
+		}
+		ln = newTcpGroupListener(group, tg, tg.lns[0].Addr())
+		realPort = tg.realPort
+		tg.lns = append(tg.lns, ln)
+	}
+	return
+}
+
+func (tg *TcpGroup) worker() {
+	for {
+		c, err := tg.tcpLn.Accept()
+		if err != nil {
+			return
+		}
+		err = gerr.PanicToError(func() {
+			tg.acceptCh <- c
+		})
+		if err != nil {
+			return
+		}
+	}
+}
+
+func (tg *TcpGroup) Accept() <-chan net.Conn {
+	return tg.acceptCh
+}
+
+func (tg *TcpGroup) CloseListener(ln *TcpGroupListener) {
+	tg.mu.Lock()
+	defer tg.mu.Unlock()
+	for i, tmpLn := range tg.lns {
+		if tmpLn == ln {
+			tg.lns = append(tg.lns[:i], tg.lns[i+1:]...)
+			break
+		}
+	}
+	if len(tg.lns) == 0 {
+		close(tg.acceptCh)
+		tg.tcpLn.Close()
+		tg.ctl.portManager.Release(tg.realPort)
+		tg.ctl.RemoveGroup(tg.group)
+	}
+}
+
+type TcpGroupCtl struct {
+	groups map[string]*TcpGroup
+
+	portManager *ports.PortManager
+	mu          sync.Mutex
+}
+
+func NewTcpGroupCtl(portManager *ports.PortManager) *TcpGroupCtl {
+	return &TcpGroupCtl{
+		groups:      make(map[string]*TcpGroup),
+		portManager: portManager,
+	}
+}
+
+func (tgc *TcpGroupCtl) Listen(proxyNanme string, group string, groupKey string,
+	addr string, port int) (l net.Listener, realPort int, err error) {
+
+	tgc.mu.Lock()
+	defer tgc.mu.Unlock()
+	if tcpGroup, ok := tgc.groups[group]; ok {
+		return tcpGroup.Listen(proxyNanme, group, groupKey, addr, port)
+	} else {
+		tcpGroup = NewTcpGroup(tgc)
+		tgc.groups[group] = tcpGroup
+		return tcpGroup.Listen(proxyNanme, group, groupKey, addr, port)
+	}
+}
+
+func (tgc *TcpGroupCtl) RemoveGroup(group string) {
+	tgc.mu.Lock()
+	defer tgc.mu.Unlock()
+	delete(tgc.groups, group)
+}