mirror of https://github.com/fatedier/frp
fatedier
8 years ago
parent
fe32a7c4bb
commit
c081df40e1
@ -0,0 +1,22 @@
|
||||
# Compiled Object files, Static and Dynamic libs (Shared Objects)
|
||||
*.o
|
||||
*.a
|
||||
*.so
|
||||
|
||||
# Folders
|
||||
_obj
|
||||
_test
|
||||
|
||||
# Architecture specific extensions/prefixes
|
||||
*.[568vq]
|
||||
[568vq].out
|
||||
|
||||
*.cgo1.go
|
||||
*.cgo2.c
|
||||
_cgo_defun.c
|
||||
_cgo_gotypes.go
|
||||
_cgo_export.*
|
||||
|
||||
_testmain.go
|
||||
|
||||
*.exe
|
||||
@ -0,0 +1,4 @@
|
||||
language: go
|
||||
go:
|
||||
- 1.1
|
||||
- tip
|
||||
@ -0,0 +1,20 @@
|
||||
The MIT License (MIT)
|
||||
|
||||
Copyright (c) 2014 Armon Dadgar
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy of
|
||||
this software and associated documentation files (the "Software"), to deal in
|
||||
the Software without restriction, including without limitation the rights to
|
||||
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
|
||||
the Software, and to permit persons to whom the Software is furnished to do so,
|
||||
subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
|
||||
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
|
||||
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
||||
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
|
||||
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
||||
@ -0,0 +1,45 @@
|
||||
go-socks5 [](https://travis-ci.org/armon/go-socks5)
|
||||
=========
|
||||
|
||||
Provides the `socks5` package that implements a [SOCKS5 server](http://en.wikipedia.org/wiki/SOCKS).
|
||||
SOCKS (Secure Sockets) is used to route traffic between a client and server through
|
||||
an intermediate proxy layer. This can be used to bypass firewalls or NATs.
|
||||
|
||||
Feature
|
||||
=======
|
||||
|
||||
The package has the following features:
|
||||
* "No Auth" mode
|
||||
* User/Password authentication
|
||||
* Support for the CONNECT command
|
||||
* Rules to do granular filtering of commands
|
||||
* Custom DNS resolution
|
||||
* Unit tests
|
||||
|
||||
TODO
|
||||
====
|
||||
|
||||
The package still needs the following:
|
||||
* Support for the BIND command
|
||||
* Support for the ASSOCIATE command
|
||||
|
||||
|
||||
Example
|
||||
=======
|
||||
|
||||
Below is a simple example of usage
|
||||
|
||||
```go
|
||||
// Create a SOCKS5 server
|
||||
conf := &socks5.Config{}
|
||||
server, err := socks5.New(conf)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
// Create SOCKS5 proxy on localhost port 8000
|
||||
if err := server.ListenAndServe("tcp", "127.0.0.1:8000"); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
```
|
||||
|
||||
@ -0,0 +1,151 @@
|
||||
package socks5
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
)
|
||||
|
||||
const (
|
||||
NoAuth = uint8(0)
|
||||
noAcceptable = uint8(255)
|
||||
UserPassAuth = uint8(2)
|
||||
userAuthVersion = uint8(1)
|
||||
authSuccess = uint8(0)
|
||||
authFailure = uint8(1)
|
||||
)
|
||||
|
||||
var (
|
||||
UserAuthFailed = fmt.Errorf("User authentication failed")
|
||||
NoSupportedAuth = fmt.Errorf("No supported authentication mechanism")
|
||||
)
|
||||
|
||||
// A Request encapsulates authentication state provided
|
||||
// during negotiation
|
||||
type AuthContext struct {
|
||||
// Provided auth method
|
||||
Method uint8
|
||||
// Payload provided during negotiation.
|
||||
// Keys depend on the used auth method.
|
||||
// For UserPassauth contains Username
|
||||
Payload map[string]string
|
||||
}
|
||||
|
||||
type Authenticator interface {
|
||||
Authenticate(reader io.Reader, writer io.Writer) (*AuthContext, error)
|
||||
GetCode() uint8
|
||||
}
|
||||
|
||||
// NoAuthAuthenticator is used to handle the "No Authentication" mode
|
||||
type NoAuthAuthenticator struct{}
|
||||
|
||||
func (a NoAuthAuthenticator) GetCode() uint8 {
|
||||
return NoAuth
|
||||
}
|
||||
|
||||
func (a NoAuthAuthenticator) Authenticate(reader io.Reader, writer io.Writer) (*AuthContext, error) {
|
||||
_, err := writer.Write([]byte{socks5Version, NoAuth})
|
||||
return &AuthContext{NoAuth, nil}, err
|
||||
}
|
||||
|
||||
// UserPassAuthenticator is used to handle username/password based
|
||||
// authentication
|
||||
type UserPassAuthenticator struct {
|
||||
Credentials CredentialStore
|
||||
}
|
||||
|
||||
func (a UserPassAuthenticator) GetCode() uint8 {
|
||||
return UserPassAuth
|
||||
}
|
||||
|
||||
func (a UserPassAuthenticator) Authenticate(reader io.Reader, writer io.Writer) (*AuthContext, error) {
|
||||
// Tell the client to use user/pass auth
|
||||
if _, err := writer.Write([]byte{socks5Version, UserPassAuth}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Get the version and username length
|
||||
header := []byte{0, 0}
|
||||
if _, err := io.ReadAtLeast(reader, header, 2); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Ensure we are compatible
|
||||
if header[0] != userAuthVersion {
|
||||
return nil, fmt.Errorf("Unsupported auth version: %v", header[0])
|
||||
}
|
||||
|
||||
// Get the user name
|
||||
userLen := int(header[1])
|
||||
user := make([]byte, userLen)
|
||||
if _, err := io.ReadAtLeast(reader, user, userLen); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Get the password length
|
||||
if _, err := reader.Read(header[:1]); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Get the password
|
||||
passLen := int(header[0])
|
||||
pass := make([]byte, passLen)
|
||||
if _, err := io.ReadAtLeast(reader, pass, passLen); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Verify the password
|
||||
if a.Credentials.Valid(string(user), string(pass)) {
|
||||
if _, err := writer.Write([]byte{userAuthVersion, authSuccess}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else {
|
||||
if _, err := writer.Write([]byte{userAuthVersion, authFailure}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, UserAuthFailed
|
||||
}
|
||||
|
||||
// Done
|
||||
return &AuthContext{UserPassAuth, map[string]string{"Username": string(user)}}, nil
|
||||
}
|
||||
|
||||
// authenticate is used to handle connection authentication
|
||||
func (s *Server) authenticate(conn io.Writer, bufConn io.Reader) (*AuthContext, error) {
|
||||
// Get the methods
|
||||
methods, err := readMethods(bufConn)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Failed to get auth methods: %v", err)
|
||||
}
|
||||
|
||||
// Select a usable method
|
||||
for _, method := range methods {
|
||||
cator, found := s.authMethods[method]
|
||||
if found {
|
||||
return cator.Authenticate(bufConn, conn)
|
||||
}
|
||||
}
|
||||
|
||||
// No usable method found
|
||||
return nil, noAcceptableAuth(conn)
|
||||
}
|
||||
|
||||
// noAcceptableAuth is used to handle when we have no eligible
|
||||
// authentication mechanism
|
||||
func noAcceptableAuth(conn io.Writer) error {
|
||||
conn.Write([]byte{socks5Version, noAcceptable})
|
||||
return NoSupportedAuth
|
||||
}
|
||||
|
||||
// readMethods is used to read the number of methods
|
||||
// and proceeding auth methods
|
||||
func readMethods(r io.Reader) ([]byte, error) {
|
||||
header := []byte{0}
|
||||
if _, err := r.Read(header); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
numMethods := int(header[0])
|
||||
methods := make([]byte, numMethods)
|
||||
_, err := io.ReadAtLeast(r, methods, numMethods)
|
||||
return methods, err
|
||||
}
|
||||
@ -0,0 +1,17 @@
|
||||
package socks5
|
||||
|
||||
// CredentialStore is used to support user/pass authentication
|
||||
type CredentialStore interface {
|
||||
Valid(user, password string) bool
|
||||
}
|
||||
|
||||
// StaticCredentials enables using a map directly as a credential store
|
||||
type StaticCredentials map[string]string
|
||||
|
||||
func (s StaticCredentials) Valid(user, password string) bool {
|
||||
pass, ok := s[user]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return password == pass
|
||||
}
|
||||
@ -0,0 +1,364 @@
|
||||
package socks5
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
const (
|
||||
ConnectCommand = uint8(1)
|
||||
BindCommand = uint8(2)
|
||||
AssociateCommand = uint8(3)
|
||||
ipv4Address = uint8(1)
|
||||
fqdnAddress = uint8(3)
|
||||
ipv6Address = uint8(4)
|
||||
)
|
||||
|
||||
const (
|
||||
successReply uint8 = iota
|
||||
serverFailure
|
||||
ruleFailure
|
||||
networkUnreachable
|
||||
hostUnreachable
|
||||
connectionRefused
|
||||
ttlExpired
|
||||
commandNotSupported
|
||||
addrTypeNotSupported
|
||||
)
|
||||
|
||||
var (
|
||||
unrecognizedAddrType = fmt.Errorf("Unrecognized address type")
|
||||
)
|
||||
|
||||
// AddressRewriter is used to rewrite a destination transparently
|
||||
type AddressRewriter interface {
|
||||
Rewrite(ctx context.Context, request *Request) (context.Context, *AddrSpec)
|
||||
}
|
||||
|
||||
// AddrSpec is used to return the target AddrSpec
|
||||
// which may be specified as IPv4, IPv6, or a FQDN
|
||||
type AddrSpec struct {
|
||||
FQDN string
|
||||
IP net.IP
|
||||
Port int
|
||||
}
|
||||
|
||||
func (a *AddrSpec) String() string {
|
||||
if a.FQDN != "" {
|
||||
return fmt.Sprintf("%s (%s):%d", a.FQDN, a.IP, a.Port)
|
||||
}
|
||||
return fmt.Sprintf("%s:%d", a.IP, a.Port)
|
||||
}
|
||||
|
||||
// Address returns a string suitable to dial; prefer returning IP-based
|
||||
// address, fallback to FQDN
|
||||
func (a AddrSpec) Address() string {
|
||||
if 0 != len(a.IP) {
|
||||
return net.JoinHostPort(a.IP.String(), strconv.Itoa(a.Port))
|
||||
}
|
||||
return net.JoinHostPort(a.FQDN, strconv.Itoa(a.Port))
|
||||
}
|
||||
|
||||
// A Request represents request received by a server
|
||||
type Request struct {
|
||||
// Protocol version
|
||||
Version uint8
|
||||
// Requested command
|
||||
Command uint8
|
||||
// AuthContext provided during negotiation
|
||||
AuthContext *AuthContext
|
||||
// AddrSpec of the the network that sent the request
|
||||
RemoteAddr *AddrSpec
|
||||
// AddrSpec of the desired destination
|
||||
DestAddr *AddrSpec
|
||||
// AddrSpec of the actual destination (might be affected by rewrite)
|
||||
realDestAddr *AddrSpec
|
||||
bufConn io.Reader
|
||||
}
|
||||
|
||||
type conn interface {
|
||||
Write([]byte) (int, error)
|
||||
RemoteAddr() net.Addr
|
||||
}
|
||||
|
||||
// NewRequest creates a new Request from the tcp connection
|
||||
func NewRequest(bufConn io.Reader) (*Request, error) {
|
||||
// Read the version byte
|
||||
header := []byte{0, 0, 0}
|
||||
if _, err := io.ReadAtLeast(bufConn, header, 3); err != nil {
|
||||
return nil, fmt.Errorf("Failed to get command version: %v", err)
|
||||
}
|
||||
|
||||
// Ensure we are compatible
|
||||
if header[0] != socks5Version {
|
||||
return nil, fmt.Errorf("Unsupported command version: %v", header[0])
|
||||
}
|
||||
|
||||
// Read in the destination address
|
||||
dest, err := readAddrSpec(bufConn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
request := &Request{
|
||||
Version: socks5Version,
|
||||
Command: header[1],
|
||||
DestAddr: dest,
|
||||
bufConn: bufConn,
|
||||
}
|
||||
|
||||
return request, nil
|
||||
}
|
||||
|
||||
// handleRequest is used for request processing after authentication
|
||||
func (s *Server) handleRequest(req *Request, conn conn) error {
|
||||
ctx := context.Background()
|
||||
|
||||
// Resolve the address if we have a FQDN
|
||||
dest := req.DestAddr
|
||||
if dest.FQDN != "" {
|
||||
ctx_, addr, err := s.config.Resolver.Resolve(ctx, dest.FQDN)
|
||||
if err != nil {
|
||||
if err := sendReply(conn, hostUnreachable, nil); err != nil {
|
||||
return fmt.Errorf("Failed to send reply: %v", err)
|
||||
}
|
||||
return fmt.Errorf("Failed to resolve destination '%v': %v", dest.FQDN, err)
|
||||
}
|
||||
ctx = ctx_
|
||||
dest.IP = addr
|
||||
}
|
||||
|
||||
// Apply any address rewrites
|
||||
req.realDestAddr = req.DestAddr
|
||||
if s.config.Rewriter != nil {
|
||||
ctx, req.realDestAddr = s.config.Rewriter.Rewrite(ctx, req)
|
||||
}
|
||||
|
||||
// Switch on the command
|
||||
switch req.Command {
|
||||
case ConnectCommand:
|
||||
return s.handleConnect(ctx, conn, req)
|
||||
case BindCommand:
|
||||
return s.handleBind(ctx, conn, req)
|
||||
case AssociateCommand:
|
||||
return s.handleAssociate(ctx, conn, req)
|
||||
default:
|
||||
if err := sendReply(conn, commandNotSupported, nil); err != nil {
|
||||
return fmt.Errorf("Failed to send reply: %v", err)
|
||||
}
|
||||
return fmt.Errorf("Unsupported command: %v", req.Command)
|
||||
}
|
||||
}
|
||||
|
||||
// handleConnect is used to handle a connect command
|
||||
func (s *Server) handleConnect(ctx context.Context, conn conn, req *Request) error {
|
||||
// Check if this is allowed
|
||||
if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
|
||||
if err := sendReply(conn, ruleFailure, nil); err != nil {
|
||||
return fmt.Errorf("Failed to send reply: %v", err)
|
||||
}
|
||||
return fmt.Errorf("Connect to %v blocked by rules", req.DestAddr)
|
||||
} else {
|
||||
ctx = ctx_
|
||||
}
|
||||
|
||||
// Attempt to connect
|
||||
dial := s.config.Dial
|
||||
if dial == nil {
|
||||
dial = func(ctx context.Context, net_, addr string) (net.Conn, error) {
|
||||
return net.Dial(net_, addr)
|
||||
}
|
||||
}
|
||||
target, err := dial(ctx, "tcp", req.realDestAddr.Address())
|
||||
if err != nil {
|
||||
msg := err.Error()
|
||||
resp := hostUnreachable
|
||||
if strings.Contains(msg, "refused") {
|
||||
resp = connectionRefused
|
||||
} else if strings.Contains(msg, "network is unreachable") {
|
||||
resp = networkUnreachable
|
||||
}
|
||||
if err := sendReply(conn, resp, nil); err != nil {
|
||||
return fmt.Errorf("Failed to send reply: %v", err)
|
||||
}
|
||||
return fmt.Errorf("Connect to %v failed: %v", req.DestAddr, err)
|
||||
}
|
||||
defer target.Close()
|
||||
|
||||
// Send success
|
||||
local := target.LocalAddr().(*net.TCPAddr)
|
||||
bind := AddrSpec{IP: local.IP, Port: local.Port}
|
||||
if err := sendReply(conn, successReply, &bind); err != nil {
|
||||
return fmt.Errorf("Failed to send reply: %v", err)
|
||||
}
|
||||
|
||||
// Start proxying
|
||||
errCh := make(chan error, 2)
|
||||
go proxy(target, req.bufConn, errCh)
|
||||
go proxy(conn, target, errCh)
|
||||
|
||||
// Wait
|
||||
for i := 0; i < 2; i++ {
|
||||
e := <-errCh
|
||||
if e != nil {
|
||||
// return from this function closes target (and conn).
|
||||
return e
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleBind is used to handle a connect command
|
||||
func (s *Server) handleBind(ctx context.Context, conn conn, req *Request) error {
|
||||
// Check if this is allowed
|
||||
if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
|
||||
if err := sendReply(conn, ruleFailure, nil); err != nil {
|
||||
return fmt.Errorf("Failed to send reply: %v", err)
|
||||
}
|
||||
return fmt.Errorf("Bind to %v blocked by rules", req.DestAddr)
|
||||
} else {
|
||||
ctx = ctx_
|
||||
}
|
||||
|
||||
// TODO: Support bind
|
||||
if err := sendReply(conn, commandNotSupported, nil); err != nil {
|
||||
return fmt.Errorf("Failed to send reply: %v", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleAssociate is used to handle a connect command
|
||||
func (s *Server) handleAssociate(ctx context.Context, conn conn, req *Request) error {
|
||||
// Check if this is allowed
|
||||
if ctx_, ok := s.config.Rules.Allow(ctx, req); !ok {
|
||||
if err := sendReply(conn, ruleFailure, nil); err != nil {
|
||||
return fmt.Errorf("Failed to send reply: %v", err)
|
||||
}
|
||||
return fmt.Errorf("Associate to %v blocked by rules", req.DestAddr)
|
||||
} else {
|
||||
ctx = ctx_
|
||||
}
|
||||
|
||||
// TODO: Support associate
|
||||
if err := sendReply(conn, commandNotSupported, nil); err != nil {
|
||||
return fmt.Errorf("Failed to send reply: %v", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// readAddrSpec is used to read AddrSpec.
|
||||
// Expects an address type byte, follwed by the address and port
|
||||
func readAddrSpec(r io.Reader) (*AddrSpec, error) {
|
||||
d := &AddrSpec{}
|
||||
|
||||
// Get the address type
|
||||
addrType := []byte{0}
|
||||
if _, err := r.Read(addrType); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Handle on a per type basis
|
||||
switch addrType[0] {
|
||||
case ipv4Address:
|
||||
addr := make([]byte, 4)
|
||||
if _, err := io.ReadAtLeast(r, addr, len(addr)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
d.IP = net.IP(addr)
|
||||
|
||||
case ipv6Address:
|
||||
addr := make([]byte, 16)
|
||||
if _, err := io.ReadAtLeast(r, addr, len(addr)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
d.IP = net.IP(addr)
|
||||
|
||||
case fqdnAddress:
|
||||
if _, err := r.Read(addrType); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
addrLen := int(addrType[0])
|
||||
fqdn := make([]byte, addrLen)
|
||||
if _, err := io.ReadAtLeast(r, fqdn, addrLen); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
d.FQDN = string(fqdn)
|
||||
|
||||
default:
|
||||
return nil, unrecognizedAddrType
|
||||
}
|
||||
|
||||
// Read the port
|
||||
port := []byte{0, 0}
|
||||
if _, err := io.ReadAtLeast(r, port, 2); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
d.Port = (int(port[0]) << 8) | int(port[1])
|
||||
|
||||
return d, nil
|
||||
}
|
||||
|
||||
// sendReply is used to send a reply message
|
||||
func sendReply(w io.Writer, resp uint8, addr *AddrSpec) error {
|
||||
// Format the address
|
||||
var addrType uint8
|
||||
var addrBody []byte
|
||||
var addrPort uint16
|
||||
switch {
|
||||
case addr == nil:
|
||||
addrType = ipv4Address
|
||||
addrBody = []byte{0, 0, 0, 0}
|
||||
addrPort = 0
|
||||
|
||||
case addr.FQDN != "":
|
||||
addrType = fqdnAddress
|
||||
addrBody = append([]byte{byte(len(addr.FQDN))}, addr.FQDN...)
|
||||
addrPort = uint16(addr.Port)
|
||||
|
||||
case addr.IP.To4() != nil:
|
||||
addrType = ipv4Address
|
||||
addrBody = []byte(addr.IP.To4())
|
||||
addrPort = uint16(addr.Port)
|
||||
|
||||
case addr.IP.To16() != nil:
|
||||
addrType = ipv6Address
|
||||
addrBody = []byte(addr.IP.To16())
|
||||
addrPort = uint16(addr.Port)
|
||||
|
||||
default:
|
||||
return fmt.Errorf("Failed to format address: %v", addr)
|
||||
}
|
||||
|
||||
// Format the message
|
||||
msg := make([]byte, 6+len(addrBody))
|
||||
msg[0] = socks5Version
|
||||
msg[1] = resp
|
||||
msg[2] = 0 // Reserved
|
||||
msg[3] = addrType
|
||||
copy(msg[4:], addrBody)
|
||||
msg[4+len(addrBody)] = byte(addrPort >> 8)
|
||||
msg[4+len(addrBody)+1] = byte(addrPort & 0xff)
|
||||
|
||||
// Send the message
|
||||
_, err := w.Write(msg)
|
||||
return err
|
||||
}
|
||||
|
||||
type closeWriter interface {
|
||||
CloseWrite() error
|
||||
}
|
||||
|
||||
// proxy is used to suffle data from src to destination, and sends errors
|
||||
// down a dedicated channel
|
||||
func proxy(dst io.Writer, src io.Reader, errCh chan error) {
|
||||
_, err := io.Copy(dst, src)
|
||||
if tcpConn, ok := dst.(closeWriter); ok {
|
||||
tcpConn.CloseWrite()
|
||||
}
|
||||
errCh <- err
|
||||
}
|
||||
@ -0,0 +1,23 @@
|
||||
package socks5
|
||||
|
||||
import (
|
||||
"net"
|
||||
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
// NameResolver is used to implement custom name resolution
|
||||
type NameResolver interface {
|
||||
Resolve(ctx context.Context, name string) (context.Context, net.IP, error)
|
||||
}
|
||||
|
||||
// DNSResolver uses the system DNS to resolve host names
|
||||
type DNSResolver struct{}
|
||||
|
||||
func (d DNSResolver) Resolve(ctx context.Context, name string) (context.Context, net.IP, error) {
|
||||
addr, err := net.ResolveIPAddr("ip", name)
|
||||
if err != nil {
|
||||
return ctx, nil, err
|
||||
}
|
||||
return ctx, addr.IP, err
|
||||
}
|
||||
@ -0,0 +1,41 @@
|
||||
package socks5
|
||||
|
||||
import (
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
// RuleSet is used to provide custom rules to allow or prohibit actions
|
||||
type RuleSet interface {
|
||||
Allow(ctx context.Context, req *Request) (context.Context, bool)
|
||||
}
|
||||
|
||||
// PermitAll returns a RuleSet which allows all types of connections
|
||||
func PermitAll() RuleSet {
|
||||
return &PermitCommand{true, true, true}
|
||||
}
|
||||
|
||||
// PermitNone returns a RuleSet which disallows all types of connections
|
||||
func PermitNone() RuleSet {
|
||||
return &PermitCommand{false, false, false}
|
||||
}
|
||||
|
||||
// PermitCommand is an implementation of the RuleSet which
|
||||
// enables filtering supported commands
|
||||
type PermitCommand struct {
|
||||
EnableConnect bool
|
||||
EnableBind bool
|
||||
EnableAssociate bool
|
||||
}
|
||||
|
||||
func (p *PermitCommand) Allow(ctx context.Context, req *Request) (context.Context, bool) {
|
||||
switch req.Command {
|
||||
case ConnectCommand:
|
||||
return ctx, p.EnableConnect
|
||||
case BindCommand:
|
||||
return ctx, p.EnableBind
|
||||
case AssociateCommand:
|
||||
return ctx, p.EnableAssociate
|
||||
}
|
||||
|
||||
return ctx, false
|
||||
}
|
||||
@ -0,0 +1,169 @@
|
||||
package socks5
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"log"
|
||||
"net"
|
||||
"os"
|
||||
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
const (
|
||||
socks5Version = uint8(5)
|
||||
)
|
||||
|
||||
// Config is used to setup and configure a Server
|
||||
type Config struct {
|
||||
// AuthMethods can be provided to implement custom authentication
|
||||
// By default, "auth-less" mode is enabled.
|
||||
// For password-based auth use UserPassAuthenticator.
|
||||
AuthMethods []Authenticator
|
||||
|
||||
// If provided, username/password authentication is enabled,
|
||||
// by appending a UserPassAuthenticator to AuthMethods. If not provided,
|
||||
// and AUthMethods is nil, then "auth-less" mode is enabled.
|
||||
Credentials CredentialStore
|
||||
|
||||
// Resolver can be provided to do custom name resolution.
|
||||
// Defaults to DNSResolver if not provided.
|
||||
Resolver NameResolver
|
||||
|
||||
// Rules is provided to enable custom logic around permitting
|
||||
// various commands. If not provided, PermitAll is used.
|
||||
Rules RuleSet
|
||||
|
||||
// Rewriter can be used to transparently rewrite addresses.
|
||||
// This is invoked before the RuleSet is invoked.
|
||||
// Defaults to NoRewrite.
|
||||
Rewriter AddressRewriter
|
||||
|
||||
// BindIP is used for bind or udp associate
|
||||
BindIP net.IP
|
||||
|
||||
// Logger can be used to provide a custom log target.
|
||||
// Defaults to stdout.
|
||||
Logger *log.Logger
|
||||
|
||||
// Optional function for dialing out
|
||||
Dial func(ctx context.Context, network, addr string) (net.Conn, error)
|
||||
}
|
||||
|
||||
// Server is reponsible for accepting connections and handling
|
||||
// the details of the SOCKS5 protocol
|
||||
type Server struct {
|
||||
config *Config
|
||||
authMethods map[uint8]Authenticator
|
||||
}
|
||||
|
||||
// New creates a new Server and potentially returns an error
|
||||
func New(conf *Config) (*Server, error) {
|
||||
// Ensure we have at least one authentication method enabled
|
||||
if len(conf.AuthMethods) == 0 {
|
||||
if conf.Credentials != nil {
|
||||
conf.AuthMethods = []Authenticator{&UserPassAuthenticator{conf.Credentials}}
|
||||
} else {
|
||||
conf.AuthMethods = []Authenticator{&NoAuthAuthenticator{}}
|
||||
}
|
||||
}
|
||||
|
||||
// Ensure we have a DNS resolver
|
||||
if conf.Resolver == nil {
|
||||
conf.Resolver = DNSResolver{}
|
||||
}
|
||||
|
||||
// Ensure we have a rule set
|
||||
if conf.Rules == nil {
|
||||
conf.Rules = PermitAll()
|
||||
}
|
||||
|
||||
// Ensure we have a log target
|
||||
if conf.Logger == nil {
|
||||
conf.Logger = log.New(os.Stdout, "", log.LstdFlags)
|
||||
}
|
||||
|
||||
server := &Server{
|
||||
config: conf,
|
||||
}
|
||||
|
||||
server.authMethods = make(map[uint8]Authenticator)
|
||||
|
||||
for _, a := range conf.AuthMethods {
|
||||
server.authMethods[a.GetCode()] = a
|
||||
}
|
||||
|
||||
return server, nil
|
||||
}
|
||||
|
||||
// ListenAndServe is used to create a listener and serve on it
|
||||
func (s *Server) ListenAndServe(network, addr string) error {
|
||||
l, err := net.Listen(network, addr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.Serve(l)
|
||||
}
|
||||
|
||||
// Serve is used to serve connections from a listener
|
||||
func (s *Server) Serve(l net.Listener) error {
|
||||
for {
|
||||
conn, err := l.Accept()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
go s.ServeConn(conn)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ServeConn is used to serve a single connection.
|
||||
func (s *Server) ServeConn(conn net.Conn) error {
|
||||
defer conn.Close()
|
||||
bufConn := bufio.NewReader(conn)
|
||||
|
||||
// Read the version byte
|
||||
version := []byte{0}
|
||||
if _, err := bufConn.Read(version); err != nil {
|
||||
s.config.Logger.Printf("[ERR] socks: Failed to get version byte: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
// Ensure we are compatible
|
||||
if version[0] != socks5Version {
|
||||
err := fmt.Errorf("Unsupported SOCKS version: %v", version)
|
||||
s.config.Logger.Printf("[ERR] socks: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
// Authenticate the connection
|
||||
authContext, err := s.authenticate(conn, bufConn)
|
||||
if err != nil {
|
||||
err = fmt.Errorf("Failed to authenticate: %v", err)
|
||||
s.config.Logger.Printf("[ERR] socks: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
request, err := NewRequest(bufConn)
|
||||
if err != nil {
|
||||
if err == unrecognizedAddrType {
|
||||
if err := sendReply(conn, addrTypeNotSupported, nil); err != nil {
|
||||
return fmt.Errorf("Failed to send reply: %v", err)
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("Failed to read destination address: %v", err)
|
||||
}
|
||||
request.AuthContext = authContext
|
||||
if client, ok := conn.RemoteAddr().(*net.TCPAddr); ok {
|
||||
request.RemoteAddr = &AddrSpec{IP: client.IP, Port: client.Port}
|
||||
}
|
||||
|
||||
// Process the client request
|
||||
if err := s.handleRequest(request, conn); err != nil {
|
||||
err = fmt.Errorf("Failed to handle request: %v", err)
|
||||
s.config.Logger.Printf("[ERR] socks: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@ -0,0 +1,156 @@
|
||||
// Copyright 2014 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// Package context defines the Context type, which carries deadlines,
|
||||
// cancelation signals, and other request-scoped values across API boundaries
|
||||
// and between processes.
|
||||
//
|
||||
// Incoming requests to a server should create a Context, and outgoing calls to
|
||||
// servers should accept a Context. The chain of function calls between must
|
||||
// propagate the Context, optionally replacing it with a modified copy created
|
||||
// using WithDeadline, WithTimeout, WithCancel, or WithValue.
|
||||
//
|
||||
// Programs that use Contexts should follow these rules to keep interfaces
|
||||
// consistent across packages and enable static analysis tools to check context
|
||||
// propagation:
|
||||
//
|
||||
// Do not store Contexts inside a struct type; instead, pass a Context
|
||||
// explicitly to each function that needs it. The Context should be the first
|
||||
// parameter, typically named ctx:
|
||||
//
|
||||
// func DoSomething(ctx context.Context, arg Arg) error {
|
||||
// // ... use ctx ...
|
||||
// }
|
||||
//
|
||||
// Do not pass a nil Context, even if a function permits it. Pass context.TODO
|
||||
// if you are unsure about which Context to use.
|
||||
//
|
||||
// Use context Values only for request-scoped data that transits processes and
|
||||
// APIs, not for passing optional parameters to functions.
|
||||
//
|
||||
// The same Context may be passed to functions running in different goroutines;
|
||||
// Contexts are safe for simultaneous use by multiple goroutines.
|
||||
//
|
||||
// See http://blog.golang.org/context for example code for a server that uses
|
||||
// Contexts.
|
||||
package context
|
||||
|
||||
import "time"
|
||||
|
||||
// A Context carries a deadline, a cancelation signal, and other values across
|
||||
// API boundaries.
|
||||
//
|
||||
// Context's methods may be called by multiple goroutines simultaneously.
|
||||
type Context interface {
|
||||
// Deadline returns the time when work done on behalf of this context
|
||||
// should be canceled. Deadline returns ok==false when no deadline is
|
||||
// set. Successive calls to Deadline return the same results.
|
||||
Deadline() (deadline time.Time, ok bool)
|
||||
|
||||
// Done returns a channel that's closed when work done on behalf of this
|
||||
// context should be canceled. Done may return nil if this context can
|
||||
// never be canceled. Successive calls to Done return the same value.
|
||||
//
|
||||
// WithCancel arranges for Done to be closed when cancel is called;
|
||||
// WithDeadline arranges for Done to be closed when the deadline
|
||||
// expires; WithTimeout arranges for Done to be closed when the timeout
|
||||
// elapses.
|
||||
//
|
||||
// Done is provided for use in select statements:
|
||||
//
|
||||
// // Stream generates values with DoSomething and sends them to out
|
||||
// // until DoSomething returns an error or ctx.Done is closed.
|
||||
// func Stream(ctx context.Context, out chan<- Value) error {
|
||||
// for {
|
||||
// v, err := DoSomething(ctx)
|
||||
// if err != nil {
|
||||
// return err
|
||||
// }
|
||||
// select {
|
||||
// case <-ctx.Done():
|
||||
// return ctx.Err()
|
||||
// case out <- v:
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// See http://blog.golang.org/pipelines for more examples of how to use
|
||||
// a Done channel for cancelation.
|
||||
Done() <-chan struct{}
|
||||
|
||||
// Err returns a non-nil error value after Done is closed. Err returns
|
||||
// Canceled if the context was canceled or DeadlineExceeded if the
|
||||
// context's deadline passed. No other values for Err are defined.
|
||||
// After Done is closed, successive calls to Err return the same value.
|
||||
Err() error
|
||||
|
||||
// Value returns the value associated with this context for key, or nil
|
||||
// if no value is associated with key. Successive calls to Value with
|
||||
// the same key returns the same result.
|
||||
//
|
||||
// Use context values only for request-scoped data that transits
|
||||
// processes and API boundaries, not for passing optional parameters to
|
||||
// functions.
|
||||
//
|
||||
// A key identifies a specific value in a Context. Functions that wish
|
||||
// to store values in Context typically allocate a key in a global
|
||||
// variable then use that key as the argument to context.WithValue and
|
||||
// Context.Value. A key can be any type that supports equality;
|
||||
// packages should define keys as an unexported type to avoid
|
||||
// collisions.
|
||||
//
|
||||
// Packages that define a Context key should provide type-safe accessors
|
||||
// for the values stores using that key:
|
||||
//
|
||||
// // Package user defines a User type that's stored in Contexts.
|
||||
// package user
|
||||
//
|
||||
// import "golang.org/x/net/context"
|
||||
//
|
||||
// // User is the type of value stored in the Contexts.
|
||||
// type User struct {...}
|
||||
//
|
||||
// // key is an unexported type for keys defined in this package.
|
||||
// // This prevents collisions with keys defined in other packages.
|
||||
// type key int
|
||||
//
|
||||
// // userKey is the key for user.User values in Contexts. It is
|
||||
// // unexported; clients use user.NewContext and user.FromContext
|
||||
// // instead of using this key directly.
|
||||
// var userKey key = 0
|
||||
//
|
||||
// // NewContext returns a new Context that carries value u.
|
||||
// func NewContext(ctx context.Context, u *User) context.Context {
|
||||
// return context.WithValue(ctx, userKey, u)
|
||||
// }
|
||||
//
|
||||
// // FromContext returns the User value stored in ctx, if any.
|
||||
// func FromContext(ctx context.Context) (*User, bool) {
|
||||
// u, ok := ctx.Value(userKey).(*User)
|
||||
// return u, ok
|
||||
// }
|
||||
Value(key interface{}) interface{}
|
||||
}
|
||||
|
||||
// Background returns a non-nil, empty Context. It is never canceled, has no
|
||||
// values, and has no deadline. It is typically used by the main function,
|
||||
// initialization, and tests, and as the top-level Context for incoming
|
||||
// requests.
|
||||
func Background() Context {
|
||||
return background
|
||||
}
|
||||
|
||||
// TODO returns a non-nil, empty Context. Code should use context.TODO when
|
||||
// it's unclear which Context to use or it is not yet available (because the
|
||||
// surrounding function has not yet been extended to accept a Context
|
||||
// parameter). TODO is recognized by static analysis tools that determine
|
||||
// whether Contexts are propagated correctly in a program.
|
||||
func TODO() Context {
|
||||
return todo
|
||||
}
|
||||
|
||||
// A CancelFunc tells an operation to abandon its work.
|
||||
// A CancelFunc does not wait for the work to stop.
|
||||
// After the first call, subsequent calls to a CancelFunc do nothing.
|
||||
type CancelFunc func()
|
||||
@ -0,0 +1,72 @@
|
||||
// Copyright 2016 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// +build go1.7
|
||||
|
||||
package context
|
||||
|
||||
import (
|
||||
"context" // standard library's context, as of Go 1.7
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
todo = context.TODO()
|
||||
background = context.Background()
|
||||
)
|
||||
|
||||
// Canceled is the error returned by Context.Err when the context is canceled.
|
||||
var Canceled = context.Canceled
|
||||
|
||||
// DeadlineExceeded is the error returned by Context.Err when the context's
|
||||
// deadline passes.
|
||||
var DeadlineExceeded = context.DeadlineExceeded
|
||||
|
||||
// WithCancel returns a copy of parent with a new Done channel. The returned
|
||||
// context's Done channel is closed when the returned cancel function is called
|
||||
// or when the parent context's Done channel is closed, whichever happens first.
|
||||
//
|
||||
// Canceling this context releases resources associated with it, so code should
|
||||
// call cancel as soon as the operations running in this Context complete.
|
||||
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
|
||||
ctx, f := context.WithCancel(parent)
|
||||
return ctx, CancelFunc(f)
|
||||
}
|
||||
|
||||
// WithDeadline returns a copy of the parent context with the deadline adjusted
|
||||
// to be no later than d. If the parent's deadline is already earlier than d,
|
||||
// WithDeadline(parent, d) is semantically equivalent to parent. The returned
|
||||
// context's Done channel is closed when the deadline expires, when the returned
|
||||
// cancel function is called, or when the parent context's Done channel is
|
||||
// closed, whichever happens first.
|
||||
//
|
||||
// Canceling this context releases resources associated with it, so code should
|
||||
// call cancel as soon as the operations running in this Context complete.
|
||||
func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc) {
|
||||
ctx, f := context.WithDeadline(parent, deadline)
|
||||
return ctx, CancelFunc(f)
|
||||
}
|
||||
|
||||
// WithTimeout returns WithDeadline(parent, time.Now().Add(timeout)).
|
||||
//
|
||||
// Canceling this context releases resources associated with it, so code should
|
||||
// call cancel as soon as the operations running in this Context complete:
|
||||
//
|
||||
// func slowOperationWithTimeout(ctx context.Context) (Result, error) {
|
||||
// ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
|
||||
// defer cancel() // releases resources if slowOperation completes before timeout elapses
|
||||
// return slowOperation(ctx)
|
||||
// }
|
||||
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
|
||||
return WithDeadline(parent, time.Now().Add(timeout))
|
||||
}
|
||||
|
||||
// WithValue returns a copy of parent in which the value associated with key is
|
||||
// val.
|
||||
//
|
||||
// Use context Values only for request-scoped data that transits processes and
|
||||
// APIs, not for passing optional parameters to functions.
|
||||
func WithValue(parent Context, key interface{}, val interface{}) Context {
|
||||
return context.WithValue(parent, key, val)
|
||||
}
|
||||
@ -0,0 +1,300 @@
|
||||
// Copyright 2014 The Go Authors. All rights reserved.
|
||||
// Use of this source code is governed by a BSD-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// +build !go1.7
|
||||
|
||||
package context
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// An emptyCtx is never canceled, has no values, and has no deadline. It is not
|
||||
// struct{}, since vars of this type must have distinct addresses.
|
||||
type emptyCtx int
|
||||
|
||||
func (*emptyCtx) Deadline() (deadline time.Time, ok bool) {
|
||||
return
|
||||
}
|
||||
|
||||
func (*emptyCtx) Done() <-chan struct{} {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (*emptyCtx) Err() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (*emptyCtx) Value(key interface{}) interface{} {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (e *emptyCtx) String() string {
|
||||
switch e {
|
||||
case background:
|
||||
return "context.Background"
|
||||
case todo:
|
||||
return "context.TODO"
|
||||
}
|
||||
return "unknown empty Context"
|
||||
}
|
||||
|
||||
var (
|
||||
background = new(emptyCtx)
|
||||
todo = new(emptyCtx)
|
||||
)
|
||||
|
||||
// Canceled is the error returned by Context.Err when the context is canceled.
|
||||
var Canceled = errors.New("context canceled")
|
||||
|
||||
// DeadlineExceeded is the error returned by Context.Err when the context's
|
||||
// deadline passes.
|
||||
var DeadlineExceeded = errors.New("context deadline exceeded")
|
||||
|
||||
// WithCancel returns a copy of parent with a new Done channel. The returned
|
||||
// context's Done channel is closed when the returned cancel function is called
|
||||
// or when the parent context's Done channel is closed, whichever happens first.
|
||||
//
|
||||
// Canceling this context releases resources associated with it, so code should
|
||||
// call cancel as soon as the operations running in this Context complete.
|
||||
func WithCancel(parent Context) (ctx Context, cancel CancelFunc) {
|
||||
c := newCancelCtx(parent)
|
||||
propagateCancel(parent, c)
|
||||
return c, func() { c.cancel(true, Canceled) }
|
||||
}
|
||||
|
||||
// newCancelCtx returns an initialized cancelCtx.
|
||||
func newCancelCtx(parent Context) *cancelCtx {
|
||||
return &cancelCtx{
|
||||
Context: parent,
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// propagateCancel arranges for child to be canceled when parent is.
|
||||
func propagateCancel(parent Context, child canceler) {
|
||||
if parent.Done() == nil {
|
||||
return // parent is never canceled
|
||||
}
|
||||
if p, ok := parentCancelCtx(parent); ok {
|
||||
p.mu.Lock()
|
||||
if p.err != nil {
|
||||
// parent has already been canceled
|
||||
child.cancel(false, p.err)
|
||||
} else {
|
||||
if p.children == nil {
|
||||
p.children = make(map[canceler]bool)
|
||||
}
|
||||
p.children[child] = true
|
||||
}
|
||||
p.mu.Unlock()
|
||||
} else {
|
||||
go func() {
|
||||
select {
|
||||
case <-parent.Done():
|
||||
child.cancel(false, parent.Err())
|
||||
case <-child.Done():
|
||||
}
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
// parentCancelCtx follows a chain of parent references until it finds a
|
||||
// *cancelCtx. This function understands how each of the concrete types in this
|
||||
// package represents its parent.
|
||||
func parentCancelCtx(parent Context) (*cancelCtx, bool) {
|
||||
for {
|
||||
switch c := parent.(type) {
|
||||
case *cancelCtx:
|
||||
return c, true
|
||||
case *timerCtx:
|
||||
return c.cancelCtx, true
|
||||
case *valueCtx:
|
||||
parent = c.Context
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// removeChild removes a context from its parent.
|
||||
func removeChild(parent Context, child canceler) {
|
||||
p, ok := parentCancelCtx(parent)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
p.mu.Lock()
|
||||
if p.children != nil {
|
||||
delete(p.children, child)
|
||||
}
|
||||
p.mu.Unlock()
|
||||
}
|
||||
|
||||
// A canceler is a context type that can be canceled directly. The
|
||||
// implementations are *cancelCtx and *timerCtx.
|
||||
type canceler interface {
|
||||
cancel(removeFromParent bool, err error)
|
||||
Done() <-chan struct{}
|
||||
}
|
||||
|
||||
// A cancelCtx can be canceled. When canceled, it also cancels any children
|
||||
// that implement canceler.
|
||||
type cancelCtx struct {
|
||||
Context
|
||||
|
||||
done chan struct{} // closed by the first cancel call.
|
||||
|
||||
mu sync.Mutex
|
||||
children map[canceler]bool // set to nil by the first cancel call
|
||||
err error // set to non-nil by the first cancel call
|
||||
}
|
||||
|
||||
func (c *cancelCtx) Done() <-chan struct{} {
|
||||
return c.done
|
||||
}
|
||||
|
||||
func (c *cancelCtx) Err() error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.err
|
||||
}
|
||||
|
||||
func (c *cancelCtx) String() string {
|
||||
return fmt.Sprintf("%v.WithCancel", c.Context)
|
||||
}
|
||||
|
||||
// cancel closes c.done, cancels each of c's children, and, if
|
||||
// removeFromParent is true, removes c from its parent's children.
|
||||
func (c *cancelCtx) cancel(removeFromParent bool, err error) {
|
||||
if err == nil {
|
||||
panic("context: internal error: missing cancel error")
|
||||
}
|
||||
c.mu.Lock()
|
||||
if c.err != nil {
|
||||
c.mu.Unlock()
|
||||
return // already canceled
|
||||
}
|
||||
c.err = err
|
||||
close(c.done)
|
||||
for child := range c.children {
|
||||
// NOTE: acquiring the child's lock while holding parent's lock.
|
||||
child.cancel(false, err)
|
||||
}
|
||||
c.children = nil
|
||||
c.mu.Unlock()
|
||||
|
||||
if removeFromParent {
|
||||
removeChild(c.Context, c)
|
||||
}
|
||||
}
|
||||
|
||||
// WithDeadline returns a copy of the parent context with the deadline adjusted
|
||||
// to be no later than d. If the parent's deadline is already earlier than d,
|
||||
// WithDeadline(parent, d) is semantically equivalent to parent. The returned
|
||||
// context's Done channel is closed when the deadline expires, when the returned
|
||||
// cancel function is called, or when the parent context's Done channel is
|
||||
// closed, whichever happens first.
|
||||
//
|
||||
// Canceling this context releases resources associated with it, so code should
|
||||
// call cancel as soon as the operations running in this Context complete.
|
||||
func WithDeadline(parent Context, deadline time.Time) (Context, CancelFunc) {
|
||||
if cur, ok := parent.Deadline(); ok && cur.Before(deadline) {
|
||||
// The current deadline is already sooner than the new one.
|
||||
return WithCancel(parent)
|
||||
}
|
||||
c := &timerCtx{
|
||||
cancelCtx: newCancelCtx(parent),
|
||||
deadline: deadline,
|
||||
}
|
||||
propagateCancel(parent, c)
|
||||
d := deadline.Sub(time.Now())
|
||||
if d <= 0 {
|
||||
c.cancel(true, DeadlineExceeded) // deadline has already passed
|
||||
return c, func() { c.cancel(true, Canceled) }
|
||||
}
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.err == nil {
|
||||
c.timer = time.AfterFunc(d, func() {
|
||||
c.cancel(true, DeadlineExceeded)
|
||||
})
|
||||
}
|
||||
return c, func() { c.cancel(true, Canceled) }
|
||||
}
|
||||
|
||||
// A timerCtx carries a timer and a deadline. It embeds a cancelCtx to
|
||||
// implement Done and Err. It implements cancel by stopping its timer then
|
||||
// delegating to cancelCtx.cancel.
|
||||
type timerCtx struct {
|
||||
*cancelCtx
|
||||
timer *time.Timer // Under cancelCtx.mu.
|
||||
|
||||
deadline time.Time
|
||||
}
|
||||
|
||||
func (c *timerCtx) Deadline() (deadline time.Time, ok bool) {
|
||||
return c.deadline, true
|
||||
}
|
||||
|
||||
func (c *timerCtx) String() string {
|
||||
return fmt.Sprintf("%v.WithDeadline(%s [%s])", c.cancelCtx.Context, c.deadline, c.deadline.Sub(time.Now()))
|
||||
}
|
||||
|
||||
func (c *timerCtx) cancel(removeFromParent bool, err error) {
|
||||
c.cancelCtx.cancel(false, err)
|
||||
if removeFromParent {
|
||||
// Remove this timerCtx from its parent cancelCtx's children.
|
||||
removeChild(c.cancelCtx.Context, c)
|
||||
}
|
||||
c.mu.Lock()
|
||||
if c.timer != nil {
|
||||
c.timer.Stop()
|
||||
c.timer = nil
|
||||
}
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
// WithTimeout returns WithDeadline(parent, time.Now().Add(timeout)).
|
||||
//
|
||||
// Canceling this context releases resources associated with it, so code should
|
||||
// call cancel as soon as the operations running in this Context complete:
|
||||
//
|
||||
// func slowOperationWithTimeout(ctx context.Context) (Result, error) {
|
||||
// ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
|
||||
// defer cancel() // releases resources if slowOperation completes before timeout elapses
|
||||
// return slowOperation(ctx)
|
||||
// }
|
||||
func WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc) {
|
||||
return WithDeadline(parent, time.Now().Add(timeout))
|
||||
}
|
||||
|
||||
// WithValue returns a copy of parent in which the value associated with key is
|
||||
// val.
|
||||
//
|
||||
// Use context Values only for request-scoped data that transits processes and
|
||||
// APIs, not for passing optional parameters to functions.
|
||||
func WithValue(parent Context, key interface{}, val interface{}) Context {
|
||||
return &valueCtx{parent, key, val}
|
||||
}
|
||||
|
||||
// A valueCtx carries a key-value pair. It implements Value for that key and
|
||||
// delegates all other calls to the embedded Context.
|
||||
type valueCtx struct {
|
||||
Context
|
||||
key, val interface{}
|
||||
}
|
||||
|
||||
func (c *valueCtx) String() string {
|
||||
return fmt.Sprintf("%v.WithValue(%#v, %#v)", c.Context, c.key, c.val)
|
||||
}
|
||||
|
||||
func (c *valueCtx) Value(key interface{}) interface{} {
|
||||
if c.key == key {
|
||||
return c.val
|
||||
}
|
||||
return c.Context.Value(key)
|
||||
}
|
||||
Loading…
Reference in new issue