1
0
mirror of https://github.com/v2fly/v2ray-core.git synced 2024-10-31 00:08:08 -04:00
v2fly/app/proxyman/mux/mux.go

433 lines
8.4 KiB
Go
Raw Normal View History

2017-02-07 15:11:47 -05:00
package mux
2017-04-08 19:43:25 -04:00
//go:generate go run $GOPATH/src/v2ray.com/core/tools/generrorgen/main.go -pkg mux -path App,Proxyman,Mux
2017-03-31 18:53:01 -04:00
import (
"context"
"sync"
"time"
2017-04-02 08:06:20 -04:00
"v2ray.com/core/app"
2017-04-02 03:48:30 -04:00
"v2ray.com/core/app/dispatcher"
"v2ray.com/core/app/log"
2017-03-31 18:53:01 -04:00
"v2ray.com/core/common/buf"
2017-04-02 03:48:30 -04:00
"v2ray.com/core/common/net"
2017-03-31 18:53:01 -04:00
"v2ray.com/core/common/signal"
"v2ray.com/core/proxy"
"v2ray.com/core/transport/ray"
)
2017-02-07 15:11:47 -05:00
2017-03-26 19:47:01 -04:00
const (
maxParallel = 8
maxTotal = 128
)
2017-04-02 03:48:30 -04:00
type manager interface {
remove(id uint16)
}
type session struct {
2017-03-31 18:53:01 -04:00
sync.Mutex
2017-04-02 03:48:30 -04:00
input ray.InputStream
output ray.OutputStream
parent manager
2017-03-31 18:53:01 -04:00
id uint16
uplinkClosed bool
downlinkClosed bool
}
2017-04-02 03:48:30 -04:00
func (s *session) closeUplink() {
2017-04-03 18:17:38 -04:00
var allDone bool
2017-03-31 18:53:01 -04:00
s.Lock()
s.uplinkClosed = true
2017-04-03 18:17:38 -04:00
allDone = s.uplinkClosed && s.downlinkClosed
2017-03-31 18:53:01 -04:00
s.Unlock()
2017-04-03 18:17:38 -04:00
if allDone {
2017-04-05 08:13:26 -04:00
s.parent.remove(s.id)
2017-04-03 18:17:38 -04:00
}
2017-03-31 18:53:01 -04:00
}
2017-04-02 03:48:30 -04:00
func (s *session) closeDownlink() {
2017-04-03 18:17:38 -04:00
var allDone bool
2017-03-31 18:53:01 -04:00
s.Lock()
s.downlinkClosed = true
2017-04-03 18:17:38 -04:00
allDone = s.uplinkClosed && s.downlinkClosed
2017-03-31 18:53:01 -04:00
s.Unlock()
2017-04-03 18:17:38 -04:00
if allDone {
2017-04-05 08:13:26 -04:00
s.parent.remove(s.id)
2017-04-03 18:17:38 -04:00
}
2017-03-31 18:53:01 -04:00
}
2017-04-02 15:30:21 -04:00
type ClientManager struct {
access sync.Mutex
clients []*Client
proxy proxy.Outbound
dialer proxy.Dialer
}
func NewClientManager(p proxy.Outbound, d proxy.Dialer) *ClientManager {
return &ClientManager{
proxy: p,
dialer: d,
}
}
func (m *ClientManager) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) error {
m.access.Lock()
defer m.access.Unlock()
for _, client := range m.clients {
if client.Dispatch(ctx, outboundRay) {
return nil
}
}
client, err := NewClient(m.proxy, m.dialer, m)
if err != nil {
2017-04-08 19:43:25 -04:00
return newError("failed to create client").Base(err)
2017-04-02 15:30:21 -04:00
}
m.clients = append(m.clients, client)
client.Dispatch(ctx, outboundRay)
return nil
}
func (m *ClientManager) onClientFinish() {
m.access.Lock()
defer m.access.Unlock()
2017-04-03 17:43:33 -04:00
if len(m.clients) < 10 {
return
}
activeClients := make([]*Client, 0, len(m.clients))
for _, client := range m.clients {
if !client.Closed() {
activeClients = append(activeClients, client)
2017-04-02 15:30:21 -04:00
}
}
2017-04-03 17:43:33 -04:00
m.clients = activeClients
2017-04-02 15:30:21 -04:00
}
2017-03-31 18:53:01 -04:00
type Client struct {
2017-04-05 08:13:26 -04:00
access sync.RWMutex
count uint16
sessions map[uint16]*session
inboundRay ray.InboundRay
ctx context.Context
cancel context.CancelFunc
manager *ClientManager
session2Remove chan uint16
2017-03-31 18:53:01 -04:00
}
2017-04-02 03:48:30 -04:00
var muxCoolDestination = net.TCPDestination(net.DomainAddress("v1.mux.cool"), net.Port(9527))
2017-03-31 18:53:01 -04:00
2017-04-02 15:30:21 -04:00
func NewClient(p proxy.Outbound, dialer proxy.Dialer, m *ClientManager) (*Client, error) {
2017-04-02 03:48:30 -04:00
ctx, cancel := context.WithCancel(context.Background())
ctx = proxy.ContextWithTarget(ctx, muxCoolDestination)
pipe := ray.NewRay(ctx)
2017-04-03 06:55:46 -04:00
go p.Process(ctx, pipe, dialer)
c := &Client{
2017-04-05 08:13:26 -04:00
sessions: make(map[uint16]*session, 256),
inboundRay: pipe,
ctx: ctx,
cancel: cancel,
manager: m,
count: 0,
session2Remove: make(chan uint16, 16),
2017-04-03 06:55:46 -04:00
}
go c.fetchOutput()
2017-04-05 08:13:26 -04:00
go c.monitor()
2017-04-03 06:55:46 -04:00
return c, nil
2017-02-07 15:11:47 -05:00
}
2017-03-31 18:53:01 -04:00
func (m *Client) remove(id uint16) {
2017-04-05 08:13:26 -04:00
select {
case m.session2Remove <- id:
default:
// Probably not gonna happen.
2017-04-02 03:48:30 -04:00
}
}
func (m *Client) Closed() bool {
select {
case <-m.ctx.Done():
return true
default:
return false
}
2017-03-26 19:47:01 -04:00
}
2017-04-05 08:13:26 -04:00
func (m *Client) monitor() {
2017-04-05 08:26:20 -04:00
defer m.manager.onClientFinish()
2017-04-05 08:13:26 -04:00
for {
select {
case <-m.ctx.Done():
m.cleanup()
return
case id := <-m.session2Remove:
m.access.Lock()
delete(m.sessions, id)
2017-04-05 09:22:21 -04:00
if len(m.sessions) == 0 {
m.cancel()
}
2017-04-05 08:13:26 -04:00
m.access.Unlock()
}
}
}
func (m *Client) cleanup() {
m.access.Lock()
defer m.access.Unlock()
2017-04-05 09:22:21 -04:00
m.inboundRay.InboundInput().Close()
m.inboundRay.InboundOutput().CloseError()
2017-04-05 08:13:26 -04:00
for _, s := range m.sessions {
s.closeUplink()
s.closeDownlink()
s.output.CloseError()
}
}
2017-04-03 06:55:46 -04:00
func fetchInput(ctx context.Context, s *session, output buf.Writer) {
2017-03-31 18:53:01 -04:00
dest, _ := proxy.TargetFromContext(ctx)
2017-04-02 07:43:24 -04:00
writer := &Writer{
2017-03-31 18:53:01 -04:00
dest: dest,
2017-04-02 03:48:30 -04:00
id: s.id,
2017-04-03 06:55:46 -04:00
writer: output,
}
defer writer.Close()
defer s.closeUplink()
2017-04-08 19:43:25 -04:00
log.Trace(newError("dispatching request to ", dest))
2017-04-03 06:55:46 -04:00
data, _ := s.input.ReadTimeout(time.Millisecond * 500)
if data != nil {
if err := writer.Write(data); err != nil {
2017-04-08 19:43:25 -04:00
log.Trace(newError("failed to write first payload").Base(err))
2017-04-03 06:55:46 -04:00
return
}
2017-03-31 18:53:01 -04:00
}
2017-04-04 04:24:38 -04:00
if err := buf.PipeUntilEOF(signal.BackgroundTimer(), s.input, writer); err != nil {
2017-04-08 19:43:25 -04:00
log.Trace(newError("failed to fetch all input").Base(err))
2017-04-02 15:30:21 -04:00
}
2017-03-31 18:53:01 -04:00
}
2017-04-05 07:46:11 -04:00
func waitForDone(ctx context.Context, s *session) {
<-ctx.Done()
s.closeUplink()
s.closeDownlink()
s.output.Close()
}
2017-04-02 03:48:30 -04:00
func (m *Client) Dispatch(ctx context.Context, outboundRay ray.OutboundRay) bool {
2017-03-31 18:53:01 -04:00
m.access.Lock()
defer m.access.Unlock()
2017-04-02 03:48:30 -04:00
if len(m.sessions) >= maxParallel {
return false
}
if m.count >= maxTotal {
return false
}
select {
case <-m.ctx.Done():
return false
default:
}
2017-03-31 18:53:01 -04:00
m.count++
id := m.count
2017-04-02 03:48:30 -04:00
s := &session{
input: outboundRay.OutboundInput(),
output: outboundRay.OutboundOutput(),
parent: m,
id: id,
2017-03-31 18:53:01 -04:00
}
2017-04-02 03:48:30 -04:00
m.sessions[id] = s
2017-04-03 06:55:46 -04:00
go fetchInput(ctx, s, m.inboundRay.InboundInput())
2017-04-05 07:46:11 -04:00
go waitForDone(ctx, s)
2017-04-02 03:48:30 -04:00
return true
2017-03-31 18:53:01 -04:00
}
2017-04-03 17:43:33 -04:00
func drain(reader *Reader) error {
for {
data, more, err := reader.Read()
if err != nil {
return err
}
data.Release()
if !more {
return nil
}
}
}
func pipe(reader *Reader, writer buf.Writer) error {
for {
data, more, err := reader.Read()
if err != nil {
return err
}
if err := writer.Write(data); err != nil {
return err
}
if !more {
return nil
}
}
}
2017-03-31 18:53:01 -04:00
func (m *Client) fetchOutput() {
2017-04-05 08:13:26 -04:00
defer m.cancel()
2017-03-31 18:53:01 -04:00
reader := NewReader(m.inboundRay.InboundOutput())
for {
meta, err := reader.ReadMetadata()
if err != nil {
2017-04-08 19:43:25 -04:00
log.Trace(newError("failed to read metadata").Base(err))
2017-03-31 18:53:01 -04:00
break
}
m.access.RLock()
2017-04-02 03:48:30 -04:00
s, found := m.sessions[meta.SessionID]
2017-03-31 18:53:01 -04:00
m.access.RUnlock()
if found && meta.SessionStatus == SessionStatusEnd {
2017-04-02 03:48:30 -04:00
s.closeDownlink()
s.output.Close()
2017-03-31 18:53:01 -04:00
}
if !meta.Option.Has(OptionData) {
continue
}
2017-04-03 17:43:33 -04:00
if found {
err = pipe(reader, s.output)
} else {
err = drain(reader)
}
if err != nil {
2017-04-08 19:43:25 -04:00
log.Trace(newError("failed to read data").Base(err))
2017-04-03 17:43:33 -04:00
break
2017-03-31 18:53:01 -04:00
}
}
2017-02-07 15:11:47 -05:00
}
2017-04-02 03:48:30 -04:00
type Server struct {
dispatcher dispatcher.Interface
}
2017-04-02 08:06:20 -04:00
func NewServer(ctx context.Context) *Server {
s := &Server{}
space := app.SpaceFromContext(ctx)
space.OnInitialize(func() error {
d := dispatcher.FromSpace(space)
if d == nil {
2017-04-08 19:43:25 -04:00
return newError("no dispatcher in space")
2017-04-02 08:06:20 -04:00
}
s.dispatcher = d
return nil
})
return s
}
2017-04-02 03:48:30 -04:00
func (s *Server) Dispatch(ctx context.Context, dest net.Destination) (ray.InboundRay, error) {
if dest != muxCoolDestination {
return s.dispatcher.Dispatch(ctx, dest)
}
ray := ray.NewRay(ctx)
2017-04-02 08:06:20 -04:00
worker := &ServerWorker{
dispatcher: s.dispatcher,
outboundRay: ray,
sessions: make(map[uint16]*session),
}
go worker.run(ctx)
2017-04-02 03:48:30 -04:00
return ray, nil
}
type ServerWorker struct {
dispatcher dispatcher.Interface
outboundRay ray.OutboundRay
sessions map[uint16]*session
access sync.RWMutex
}
func (w *ServerWorker) remove(id uint16) {
w.access.Lock()
delete(w.sessions, id)
w.access.Unlock()
}
2017-04-03 11:28:15 -04:00
func handle(ctx context.Context, s *session, output buf.Writer) {
writer := NewResponseWriter(s.id, output)
2017-04-04 04:24:38 -04:00
if err := buf.PipeUntilEOF(signal.BackgroundTimer(), s.input, writer); err != nil {
2017-04-08 19:43:25 -04:00
log.Trace(newError("session ", s.id, " ends: ").Base(err))
2017-04-02 03:48:30 -04:00
}
2017-04-04 05:20:07 -04:00
writer.Close()
s.closeDownlink()
2017-04-02 03:48:30 -04:00
}
func (w *ServerWorker) run(ctx context.Context) {
input := w.outboundRay.OutboundInput()
reader := NewReader(input)
for {
select {
case <-ctx.Done():
return
default:
}
meta, err := reader.ReadMetadata()
if err != nil {
2017-04-08 19:43:25 -04:00
log.Trace(newError("failed to read metadata").Base(err))
2017-04-02 03:48:30 -04:00
return
}
w.access.RLock()
s, found := w.sessions[meta.SessionID]
w.access.RUnlock()
if found && meta.SessionStatus == SessionStatusEnd {
s.closeUplink()
s.output.Close()
}
if meta.SessionStatus == SessionStatusNew {
2017-04-08 19:43:25 -04:00
log.Trace(newError("received request for ", meta.Target))
2017-04-02 03:48:30 -04:00
inboundRay, err := w.dispatcher.Dispatch(ctx, meta.Target)
if err != nil {
2017-04-08 19:43:25 -04:00
log.Trace(newError("failed to dispatch request.").Base(err))
2017-04-02 03:48:30 -04:00
continue
}
s = &session{
input: inboundRay.InboundOutput(),
output: inboundRay.InboundInput(),
parent: w,
id: meta.SessionID,
}
w.access.Lock()
w.sessions[meta.SessionID] = s
w.access.Unlock()
2017-04-03 11:28:15 -04:00
go handle(ctx, s, w.outboundRay.OutboundOutput())
2017-04-02 03:48:30 -04:00
}
2017-04-03 17:43:33 -04:00
if !meta.Option.Has(OptionData) {
continue
}
if s != nil {
err = pipe(reader, s.output)
} else {
err = drain(reader)
}
if err != nil {
2017-04-08 19:43:25 -04:00
log.Trace(newError("failed to read data").Base(err))
2017-04-03 17:43:33 -04:00
break
2017-04-02 03:48:30 -04:00
}
}
}