apply policy in mtproto inbound

pull/1524/head^2
Darien Raymond 2018-07-04 18:58:19 +02:00
parent 5c2bf477bb
commit d7def8bf47
No known key found for this signature in database
GPG Key ID: 7251FFA14BB18169
1 changed files with 25 additions and 0 deletions

View File

@ -2,6 +2,7 @@ package mtproto
import ( import (
"context" "context"
"time"
"v2ray.com/core" "v2ray.com/core"
"v2ray.com/core/common" "v2ray.com/core/common"
@ -10,6 +11,8 @@ import (
"v2ray.com/core/common/net" "v2ray.com/core/common/net"
"v2ray.com/core/common/predicate" "v2ray.com/core/common/predicate"
"v2ray.com/core/common/protocol" "v2ray.com/core/common/protocol"
"v2ray.com/core/common/session"
"v2ray.com/core/common/signal"
"v2ray.com/core/common/task" "v2ray.com/core/common/task"
"v2ray.com/core/transport/internet" "v2ray.com/core/transport/internet"
"v2ray.com/core/transport/pipe" "v2ray.com/core/transport/pipe"
@ -28,6 +31,7 @@ var (
type Server struct { type Server struct {
user *protocol.User user *protocol.User
account *Account account *Account
policy core.PolicyManager
} }
func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) { func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) {
@ -45,9 +49,12 @@ func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) {
return nil, newError("not a MTProto account") return nil, newError("not a MTProto account")
} }
v := core.MustFromContext(ctx)
return &Server{ return &Server{
user: user, user: user,
account: account, account: account,
policy: v.PolicyManager(),
}, nil }, nil
} }
@ -58,12 +65,21 @@ func (s *Server) Network() net.NetworkList {
} }
func (s *Server) Process(ctx context.Context, network net.Network, conn internet.Connection, dispatcher core.Dispatcher) error { func (s *Server) Process(ctx context.Context, network net.Network, conn internet.Connection, dispatcher core.Dispatcher) error {
sPolicy := s.policy.ForLevel(s.user.Level)
if err := conn.SetDeadline(time.Now().Add(sPolicy.Timeouts.Handshake)); err != nil {
newError("failed to set deadline").Base(err).WriteToLog(session.ExportIDToError(ctx))
}
auth, err := ReadAuthentication(conn) auth, err := ReadAuthentication(conn)
if err != nil { if err != nil {
return newError("failed to read authentication header").Base(err) return newError("failed to read authentication header").Base(err)
} }
defer putAuthenticationObject(auth) defer putAuthenticationObject(auth)
if err := conn.SetDeadline(time.Time{}); err != nil {
newError("failed to clear deadline").Base(err).WriteToLog(session.ExportIDToError(ctx))
}
auth.ApplySecret(s.account.Secret) auth.ApplySecret(s.account.Secret)
decryptor := crypto.NewAesCTRStream(auth.DecodingKey[:], auth.DecodingNonce[:]) decryptor := crypto.NewAesCTRStream(auth.DecodingKey[:], auth.DecodingNonce[:])
@ -83,17 +99,26 @@ func (s *Server) Process(ctx context.Context, network net.Network, conn internet
Address: dcList[dcID], Address: dcList[dcID],
Port: net.Port(443), Port: net.Port(443),
} }
ctx, cancel := context.WithCancel(ctx)
timer := signal.CancelAfterInactivity(ctx, cancel, sPolicy.Timeouts.ConnectionIdle)
ctx = core.ContextWithBufferPolicy(ctx, sPolicy.Buffer)
link, err := dispatcher.Dispatch(ctx, dest) link, err := dispatcher.Dispatch(ctx, dest)
if err != nil { if err != nil {
return newError("failed to dispatch request to: ", dest).Base(err) return newError("failed to dispatch request to: ", dest).Base(err)
} }
request := func() error { request := func() error {
defer timer.SetTimeout(sPolicy.Timeouts.DownlinkOnly)
reader := buf.NewReader(crypto.NewCryptionReader(decryptor, conn)) reader := buf.NewReader(crypto.NewCryptionReader(decryptor, conn))
return buf.Copy(reader, link.Writer) return buf.Copy(reader, link.Writer)
} }
response := func() error { response := func() error {
defer timer.SetTimeout(sPolicy.Timeouts.UplinkOnly)
encryptor := crypto.NewAesCTRStream(auth.EncodingKey[:], auth.EncodingNonce[:]) encryptor := crypto.NewAesCTRStream(auth.EncodingKey[:], auth.EncodingNonce[:])
writer := buf.NewWriter(crypto.NewCryptionWriter(encryptor, conn)) writer := buf.NewWriter(crypto.NewCryptionWriter(encryptor, conn))
return buf.Copy(link.Reader, writer) return buf.Copy(link.Reader, writer)