<commit_before>// R.I.P Shadowsocks
package shadowsocks

import (
	"sync"

	"errors"
	"v2ray.com/core/app"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/common"
	"v2ray.com/core/common/alloc"
	v2io "v2ray.com/core/common/io"
	"v2ray.com/core/common/log"
	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/internet/udp"
)

type Server struct {
	packetDispatcher dispatcher.PacketDispatcher
	config           *ServerConfig
	user             *protocol.User
	account          *ShadowsocksAccount
	meta             *proxy.InboundHandlerMeta
	accepting        bool
	tcpHub           *internet.TCPHub
	udpHub           *udp.UDPHub
	udpServer        *udp.UDPServer
}

func NewServer(config *ServerConfig, space app.Space, meta *proxy.InboundHandlerMeta) (*Server, error) {
	if config.GetUser() == nil {
		return nil, protocol.ErrUserMissing
	}

	rawAccount, err := config.User.GetTypedAccount()
	if err != nil {
		return nil, errors.New("Shadowsocks|Server: Failed to get user account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	s := &Server{
		config:  config,
		meta:    meta,
		user:    config.GetUser(),
		account: account,
	}

	space.InitializeApplication(func() error {
		if !space.HasApp(dispatcher.APP_ID) {
			return app.ErrMissingApplication
		}
		s.packetDispatcher = space.GetApp(dispatcher.APP_ID).(dispatcher.PacketDispatcher)
		return nil
	})

	return s, nil
}

func (this *Server) Port() v2net.Port {
	return this.meta.Port
}

func (this *Server) Close() {
	this.accepting = false
	// TODO: synchronization
	if this.tcpHub != nil {
		this.tcpHub.Close()
		this.tcpHub = nil
	}

	if this.udpHub != nil {
		this.udpHub.Close()
		this.udpHub = nil
	}
}

func (this *Server) Start() error {
	if this.accepting {
		return nil
	}

	tcpHub, err := internet.ListenTCP(this.meta.Address, this.meta.Port, this.handleConnection, this.meta.StreamSettings)
	if err != nil {
		log.Error("Shadowsocks: Failed to listen TCP on ", this.meta.Address, ":", this.meta.Port, ": ", err)
		return err
	}
	this.tcpHub = tcpHub

	if this.config.UdpEnabled {
		this.udpServer = udp.NewUDPServer(this.meta, this.packetDispatcher)
		udpHub, err := udp.ListenUDP(this.meta.Address, this.meta.Port, udp.ListenOption{Callback: this.handlerUDPPayload})
		if err != nil {
			log.Error("Shadowsocks: Failed to listen UDP on ", this.meta.Address, ":", this.meta.Port, ": ", err)
			return err
		}
		this.udpHub = udpHub
	}

	this.accepting = true

	return nil
}

func (this *Server) handlerUDPPayload(payload *alloc.Buffer, session *proxy.SessionInfo) {
	source := session.Source
	request, data, err := DecodeUDPPacket(this.user, payload)
	if err != nil {
		log.Info("Shadowsocks|Server: Skipping invalid UDP packet from: ", source, ": ", err)
		log.Access(source, "", log.AccessRejected, err)
		payload.Release()
		return
	}

	if request.Option.Has(RequestOptionOneTimeAuth) && this.account.OneTimeAuth == Account_Disabled {
		log.Info("Shadowsocks|Server: Client payload enables OTA but server doesn't allow it.")
		payload.Release()
		return
	}

	if !request.Option.Has(RequestOptionOneTimeAuth) && this.account.OneTimeAuth == Account_Enabled {
		log.Info("Shadowsocks|Server: Client payload disables OTA but server forces it.")
		payload.Release()
		return
	}

	dest := request.Destination()
	log.Access(source, dest, log.AccessAccepted, "")
	log.Info("Shadowsocks|Server: Tunnelling request to ", dest)

	this.udpServer.Dispatch(&proxy.SessionInfo{Source: source, Destination: dest, User: request.User}, data, func(destination v2net.Destination, payload *alloc.Buffer) {
		defer payload.Release()

		data, err := EncodeUDPPacket(request, payload)
		if err != nil {
			log.Warning("Shadowsocks|Server: Failed to encode UDP packet: ", err)
			return
		}
		defer data.Release()

		this.udpHub.WriteTo(data.Value, source)
	})
}

func (this *Server) handleConnection(conn internet.Connection) {
	defer conn.Close()
	conn.SetReusable(false)

	timedReader := v2net.NewTimeOutReader(16, conn)
	defer timedReader.Release()

	bufferedReader := v2io.NewBufferedReader(timedReader)
	defer bufferedReader.Release()

	request, bodyReader, err := ReadTCPSession(this.user, bufferedReader)
	if err != nil {
		log.Access(conn.RemoteAddr(), "", log.AccessRejected, err)
		log.Info("Shadowsocks|Server: Failed to create request from: ", conn.RemoteAddr(), ": ", err)
		return
	}
	defer bodyReader.Release()

	bufferedReader.SetCached(false)

	userSettings := this.user.GetSettings()
	timedReader.SetTimeOut(userSettings.PayloadReadTimeout)

	dest := request.Destination()
	log.Access(conn.RemoteAddr(), dest, log.AccessAccepted, "")
	log.Info("Shadowsocks|Server: Tunnelling request to ", dest)

	ray := this.packetDispatcher.DispatchToOutbound(this.meta, &proxy.SessionInfo{
		Source:      v2net.DestinationFromAddr(conn.RemoteAddr()),
		Destination: dest,
		User:        request.User,
	})
	defer ray.InboundOutput().Release()

	var writeFinish sync.Mutex
	writeFinish.Lock()
	go func() {
		defer writeFinish.Unlock()

		bufferedWriter := v2io.NewBufferedWriter(conn)
		defer bufferedWriter.Release()

		responseWriter, err := WriteTCPResponse(request, bufferedWriter)
		if err != nil {
			log.Warning("Shadowsocks|Server: Failed to write response: ", err)
			return
		}
		defer responseWriter.Release()

		if payload, err := ray.InboundOutput().Read(); err == nil {
			responseWriter.Write(payload)
			bufferedWriter.SetCached(false)

			v2io.Pipe(ray.InboundOutput(), responseWriter)
		}
	}()

	v2io.Pipe(bodyReader, ray.InboundInput())
	ray.InboundInput().Close()

	writeFinish.Lock()
}

type ServerFactory struct{}

func (this *ServerFactory) StreamCapability() v2net.NetworkList {
	return v2net.NetworkList{
		Network: []v2net.Network{v2net.Network_RawTCP},
	}
}

func (this *ServerFactory) Create(space app.Space, rawConfig interface{}, meta *proxy.InboundHandlerMeta) (proxy.InboundHandler, error) {
	if !space.HasApp(dispatcher.APP_ID) {
		return nil, common.ErrBadConfiguration
	}
	return NewServer(rawConfig.(*ServerConfig), space, meta)
}
<commit_msg>enable tcp for shadowsocks server<commit_after>// R.I.P Shadowsocks
package shadowsocks

import (
	"sync"

	"errors"
	"v2ray.com/core/app"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/common"
	"v2ray.com/core/common/alloc"
	v2io "v2ray.com/core/common/io"
	"v2ray.com/core/common/log"
	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/internet/udp"
)

type Server struct {
	packetDispatcher dispatcher.PacketDispatcher
	config           *ServerConfig
	user             *protocol.User
	account          *ShadowsocksAccount
	meta             *proxy.InboundHandlerMeta
	accepting        bool
	tcpHub           *internet.TCPHub
	udpHub           *udp.UDPHub
	udpServer        *udp.UDPServer
}

func NewServer(config *ServerConfig, space app.Space, meta *proxy.InboundHandlerMeta) (*Server, error) {
	if config.GetUser() == nil {
		return nil, protocol.ErrUserMissing
	}

	rawAccount, err := config.User.GetTypedAccount()
	if err != nil {
		return nil, errors.New("Shadowsocks|Server: Failed to get user account: " + err.Error())
	}
	account := rawAccount.(*ShadowsocksAccount)

	s := &Server{
		config:  config,
		meta:    meta,
		user:    config.GetUser(),
		account: account,
	}

	space.InitializeApplication(func() error {
		if !space.HasApp(dispatcher.APP_ID) {
			return app.ErrMissingApplication
		}
		s.packetDispatcher = space.GetApp(dispatcher.APP_ID).(dispatcher.PacketDispatcher)
		return nil
	})

	return s, nil
}

func (this *Server) Port() v2net.Port {
	return this.meta.Port
}

func (this *Server) Close() {
	this.accepting = false
	// TODO: synchronization
	if this.tcpHub != nil {
		this.tcpHub.Close()
		this.tcpHub = nil
	}

	if this.udpHub != nil {
		this.udpHub.Close()
		this.udpHub = nil
	}
}

func (this *Server) Start() error {
	if this.accepting {
		return nil
	}

	tcpHub, err := internet.ListenTCP(this.meta.Address, this.meta.Port, this.handleConnection, this.meta.StreamSettings)
	if err != nil {
		log.Error("Shadowsocks: Failed to listen TCP on ", this.meta.Address, ":", this.meta.Port, ": ", err)
		return err
	}
	this.tcpHub = tcpHub

	if this.config.UdpEnabled {
		this.udpServer = udp.NewUDPServer(this.meta, this.packetDispatcher)
		udpHub, err := udp.ListenUDP(this.meta.Address, this.meta.Port, udp.ListenOption{Callback: this.handlerUDPPayload})
		if err != nil {
			log.Error("Shadowsocks: Failed to listen UDP on ", this.meta.Address, ":", this.meta.Port, ": ", err)
			return err
		}
		this.udpHub = udpHub
	}

	this.accepting = true

	return nil
}

func (this *Server) handlerUDPPayload(payload *alloc.Buffer, session *proxy.SessionInfo) {
	source := session.Source
	request, data, err := DecodeUDPPacket(this.user, payload)
	if err != nil {
		log.Info("Shadowsocks|Server: Skipping invalid UDP packet from: ", source, ": ", err)
		log.Access(source, "", log.AccessRejected, err)
		payload.Release()
		return
	}

	if request.Option.Has(RequestOptionOneTimeAuth) && this.account.OneTimeAuth == Account_Disabled {
		log.Info("Shadowsocks|Server: Client payload enables OTA but server doesn't allow it.")
		payload.Release()
		return
	}

	if !request.Option.Has(RequestOptionOneTimeAuth) && this.account.OneTimeAuth == Account_Enabled {
		log.Info("Shadowsocks|Server: Client payload disables OTA but server forces it.")
		payload.Release()
		return
	}

	dest := request.Destination()
	log.Access(source, dest, log.AccessAccepted, "")
	log.Info("Shadowsocks|Server: Tunnelling request to ", dest)

	this.udpServer.Dispatch(&proxy.SessionInfo{Source: source, Destination: dest, User: request.User}, data, func(destination v2net.Destination, payload *alloc.Buffer) {
		defer payload.Release()

		data, err := EncodeUDPPacket(request, payload)
		if err != nil {
			log.Warning("Shadowsocks|Server: Failed to encode UDP packet: ", err)
			return
		}
		defer data.Release()

		this.udpHub.WriteTo(data.Value, source)
	})
}

func (this *Server) handleConnection(conn internet.Connection) {
	defer conn.Close()
	conn.SetReusable(false)

	timedReader := v2net.NewTimeOutReader(16, conn)
	defer timedReader.Release()

	bufferedReader := v2io.NewBufferedReader(timedReader)
	defer bufferedReader.Release()

	request, bodyReader, err := ReadTCPSession(this.user, bufferedReader)
	if err != nil {
		log.Access(conn.RemoteAddr(), "", log.AccessRejected, err)
		log.Info("Shadowsocks|Server: Failed to create request from: ", conn.RemoteAddr(), ": ", err)
		return
	}
	defer bodyReader.Release()

	bufferedReader.SetCached(false)

	userSettings := this.user.GetSettings()
	timedReader.SetTimeOut(userSettings.PayloadReadTimeout)

	dest := request.Destination()
	log.Access(conn.RemoteAddr(), dest, log.AccessAccepted, "")
	log.Info("Shadowsocks|Server: Tunnelling request to ", dest)

	ray := this.packetDispatcher.DispatchToOutbound(this.meta, &proxy.SessionInfo{
		Source:      v2net.DestinationFromAddr(conn.RemoteAddr()),
		Destination: dest,
		User:        request.User,
	})
	defer ray.InboundOutput().Release()

	var writeFinish sync.Mutex
	writeFinish.Lock()
	go func() {
		defer writeFinish.Unlock()

		bufferedWriter := v2io.NewBufferedWriter(conn)
		defer bufferedWriter.Release()

		responseWriter, err := WriteTCPResponse(request, bufferedWriter)
		if err != nil {
			log.Warning("Shadowsocks|Server: Failed to write response: ", err)
			return
		}
		defer responseWriter.Release()

		if payload, err := ray.InboundOutput().Read(); err == nil {
			responseWriter.Write(payload)
			bufferedWriter.SetCached(false)

			v2io.Pipe(ray.InboundOutput(), responseWriter)
		}
	}()

	v2io.Pipe(bodyReader, ray.InboundInput())
	ray.InboundInput().Close()

	writeFinish.Lock()
}

type ServerFactory struct{}

func (this *ServerFactory) StreamCapability() v2net.NetworkList {
	return v2net.NetworkList{
		Network: []v2net.Network{v2net.Network_TCP, v2net.Network_RawTCP},
	}
}

func (this *ServerFactory) Create(space app.Space, rawConfig interface{}, meta *proxy.InboundHandlerMeta) (proxy.InboundHandler, error) {
	if !space.HasApp(dispatcher.APP_ID) {
		return nil, common.ErrBadConfiguration
	}
	return NewServer(rawConfig.(*ServerConfig), space, meta)
}
<|endoftext|>


<commit_before>// nflag is an open source alternative to Golang's flag package.
// nflag allows OS-specific flag styling, type option simplification, and allowance of empty values.

package nflag

import (
	"errors"
	"fmt"
	"os"
	"reflect"
	"runtime"
	"strconv"
	"strings"
)

var Config ConfigOptions
var Flags map[string]Flag

// Package Init
func init() {
	Flags = make(map[string]Flag)      // Make the Flags map
	Config.OSSpecificFlags = true      // Default to being true
	Config.OSSpecificFlagString = "--" // Default to using a double dash for declaring flags

	if runtime.GOOS == "windows" { // If we are on Windows
		Config.OSSpecificFlagString = "/" // Use a slash, since Windows likes to be "special"
	}
}

// Configure
// This function is for configuration of nflag prior to usage.
func Configure(providedConfig ConfigOptions) {
	if providedConfig.OSSpecificFlags == false { // If we are overriding OSSpecificFlags
		Config.OSSpecificFlags = false
		Config.OSSpecificFlagString = "--" // Enforce --
	}

	if providedConfig.OSSpecificFlagString != "" { // If we are overriding OSSpecificFlagString
		Config.OSSpecificFlagString = providedConfig.OSSpecificFlagString // Set OSSpecificFlagString to whatever is provided on config
	}

	if providedConfig.ShowHelpIfNoArgs { // If ShowHelpIfNoArgs was set to true
		Config.ShowHelpIfNoArgs = true // Change to true
	}
}

// Get
// This function will get the flag value and returns it, or an error if the flag does not exist.
func Get(flagName string) (interface{}, error) {
	flag, exists := Flags[flagName] // Get the trimmedFlag and the exists bool of this flagName in Flags
	var val interface{}
	var err error

	if exists { // If the flag exists
		val = flag.Value
	} else {
		err = errors.New("Flag does not exist.")
	}

	return val, err
}

// #region GetAs functions

// GetAsBool
// This function will get the flag value and convert it to bool, or an error if the flag does not exist.
func GetAsBool(flagName string) (bool, error) {
	var flagValue bool
	flagValueInterface, flagValueError := Get(flagName)

	if flagValueError == nil { // If this flag does exist
		flagValue = flagValueInterface.(bool) // Type inference to bool
	}

	return flagValue, flagValueError
}

// GetAsInt
// This function will get the flag value and convert it to int, or an error if the flag does not exist.
func GetAsInt(flagName string) (int, error) {
	var flagValue int
	flagValueInterface, flagValueError := Get(flagName)

	if flagValueError == nil { // If this flag does exist
		flagValue = flagValueInterface.(int) // Type inference to int
	}

	return flagValue, flagValueError
}

// GetAsFloat64
// This function will get the flag value and convert it to float64, or an error if the flag does not exist.
func GetAsFloat64(flagName string) (float64, error) {
	var flagValue float64
	flagValueInterface, flagValueError := Get(flagName)

	if flagValueError == nil { // If this flag does exist
		flagValue = flagValueInterface.(float64) // Type inference to float64
	}

	return flagValue, flagValueError
}

// GetAsString
// This function will get the flag value and convert it to string, or an error if the flag does not exist.
func GetAsString(flagName string) (string, error) {
	var flagValue string
	flagValueInterface, flagValueError := Get(flagName)

	if flagValueError == nil { // If this flag does exist
		flagValue = flagValueInterface.(string) // Type inference to string
	}

	return flagValue, flagValueError
}

// #endregion

// Set
// This function is for setting a flag.
func Set(flagName string, providedFlag Flag) error {
	var errorResponse error

	if providedFlag.Type == "" { // If no Type was provided
		providedFlag.Type = "bool" // Default to bool
	}

	allowedTypes := []string{"bool", "float64", "int", "string"} // Define allowed types
	isAllowedType := false                                       // Default to providedFlag.Type not being the type

	for _, val := range allowedTypes {
		if providedFlag.Type == val { // If the Type is the same as what is allowed (this specific value)
			isAllowedType = true
		}
	}

	if isAllowedType { // If this is an allowed type
		if (providedFlag.DefaultValue == nil) && (providedFlag.AllowNothing) { // If no default value was provided and we are allowing nothing
			if providedFlag.Type == "bool" { // If the type is bool
				providedFlag.DefaultValue = false
			} else if providedFlag.Type == "float64" { // If the type is float64
				providedFlag.DefaultValue = 1.1
			} else if providedFlag.Type == "int" { // If the type is int
				providedFlag.DefaultValue = 1
			} else if providedFlag.Type == "string" {
				providedFlag.DefaultValue = ""
			}
		} else if (providedFlag.DefaultValue != nil) && (providedFlag.AllowNothing == false) { // If a Default value was provided
			providedFlag.AllowNothing = true                                             // Enforce AllowNothing since a default value is provided
			if providedFlag.Type != reflect.TypeOf(providedFlag.DefaultValue).String() { // If the Type and DefaultValue aren't the same type
				errorResponse = errors.New("Mismatch Flag and DefaultValue types.")
			}
		}

		if errorResponse == nil { // If there was no error
			Flags[flagName] = providedFlag // Set the flag name in Flags to the providedFlag struct
		}
	} else { // If this is a non-allowed type
		errorResponse = errors.New("Type is not allowed. Please use bool, float64, int, or string.")
	}

	return errorResponse
}

// Parse
// This function will parse input for flags
func Parse() {
	providedFlags := os.Args[1:] // ProvidedFlags start at index 1 (after binary name)

	if len(providedFlags) != 0 { // If flags were provided
		if providedFlags[0] != (Config.OSSpecificFlagString + "help") { // If the first providedFlag is not help
			for _, flag := range providedFlags { // For each flag provided
				flagNameValueSplit := strings.Split(flag, "=")                                          // Split the flag name and the value
				flagName := strings.Replace(flagNameValueSplit[0], Config.OSSpecificFlagString, "", -1) // Get the flagName and remove the OSSpecificFlagString
				flagValue := ""                                                                         // Default to nothing

				if len(flagNameValueSplit) == 2 { // If a value was provided
					flagValue = flagNameValueSplit[1] // Set to the value in flagNameValueSplit
				} else { // If the flag was passed but no value
					trimmedFlag, exists := Flags[flagName] // Get the trimmedFlag and the exists bool of this flagName in Flags

					if exists && (trimmedFlag.Type == "bool") { // If the flag exists and its type is bool
						flagValue = "true" // By passing just --flag, it implies --flag is true
					}
				}

				ParseVal(flagName, flagValue) // Parse the value
			}
		} else { // If we should output the help message
			PrintFlags() // Print the flags
			os.Exit(1)
		}

		for flagName, _ := range Flags { // For each flagName in Flags
			if Flags[flagName].Value == nil { // If no value is set for this Flag (which happens if it wasn't parsed via
				ParseVal(flagName, "flag-not-provided")
			}
		}
	} else if Config.ShowHelpIfNoArgs { // If no arguments are passed and ShowHelpIfNoArgs is true
		PrintFlags()
		os.Exit(1)
	}
}

// ParseVal
// This function will check Flags for the provided flag and parse the provided value
func ParseVal(flagName string, flagValue string) {
	trimmedFlag, exists := Flags[flagName] // Get the trimmedFlag and the exists bool of this flagName in Flags

	if exists { // If the flag exists
		if (flagValue != "") && (flagValue != "flag-not-provided") { // If a value was defined
			var conversionError error // Define conversionError as any error from not correctly converting the value to the type

			if trimmedFlag.Type == "bool" { // If the type is bool
				trimmedFlag.Value, conversionError = strconv.ParseBool(flagValue) // Convert to bool
			} else if trimmedFlag.Type == "float64" { // If the type is float64
				trimmedFlag.Value, conversionError = strconv.ParseFloat(flagValue, 64) // Convert to float64
			} else if trimmedFlag.Type == "int" { // If the type is int
				trimmedFlag.Value, conversionError = strconv.Atoi(flagValue) // Convert to int
			} else { // Remaining type, being int
				trimmedFlag.Value = flagValue
			}

			if conversionError != nil { // If there was a conversionError
				fmt.Println("An incorrect value was provided when using " + Config.OSSpecificFlagString + "flagName . Exiting.")
				os.Exit(1)
			}
		} else { // If no value was provided by the user
			returnRequiredValueMessage := false

			if flagValue == "" { // If the flag was provided by the user but no content
				if trimmedFlag.AllowNothing { // If providing no content is permitted
					if trimmedFlag.Type == "bool" { // If the flag is bool
						trimmedFlag.Value = true // Force true
					} else { // If the flag is not bool
						trimmedFlag.Value = trimmedFlag.DefaultValue // Set to defaultvalue
					}
				} else { // If providing no content is NOT permitted
					returnRequiredValueMessage = true
				}
			} else { // If this flag was not provided
				if !trimmedFlag.Required { // If the flag is NOT required
					trimmedFlag.Value = trimmedFlag.DefaultValue // Set the value to the default
				} else { // If the flag IS required
					returnRequiredValueMessage = true
				}
			}

			if returnRequiredValueMessage {
				fmt.Println("A required value for " + Config.OSSpecificFlagString + flagName + " was not provided. Exiting.")
				os.Exit(1)
			}
		}

		Flags[flagName] = trimmedFlag // Update the map
	} else { // If the flag does not exist.
		fmt.Println(Config.OSSpecificFlagString + flagName + " does not exist. Exiting.")
		os.Exit(1)
	}
}

// PrintFlags
// This function will print all the flags that are set and their defaults
func PrintFlags() {
	fmt.Println("Usage: " + Config.OSSpecificFlagString + "example=value")
	fmt.Println("The following options are available:")

	for flagName, flag := range Flags { // For each flagName and trimmedFlag struct in Flags
		fmt.Println(Config.OSSpecificFlagString + flagName + " " + flag.Type)
		fmt.Println("\t" + flag.Descriptor)

		if flag.DefaultValue != nil { // If DefaultValue is not nil
			fmt.Println("\t" + "Default Value: " + reflect.ValueOf(flag.DefaultValue).String())
		}

		fmt.Println("\t" + "Allows Providing Only Flag: " + strconv.FormatBool(flag.AllowNothing))
		fmt.Println("")
	}
}
<commit_msg>Assign a default value if DefaultValue is nil, and remove requirement for AllowNothing val to be true.<commit_after>// nflag is an open source alternative to Golang's flag package.
// nflag allows OS-specific flag styling, type option simplification, and allowance of empty values.

package nflag

import (
	"errors"
	"fmt"
	"os"
	"reflect"
	"runtime"
	"strconv"
	"strings"
)

var Config ConfigOptions
var Flags map[string]Flag

// Package Init
func init() {
	Flags = make(map[string]Flag)      // Make the Flags map
	Config.OSSpecificFlags = true      // Default to being true
	Config.OSSpecificFlagString = "--" // Default to using a double dash for declaring flags

	if runtime.GOOS == "windows" { // If we are on Windows
		Config.OSSpecificFlagString = "/" // Use a slash, since Windows likes to be "special"
	}
}

// Configure
// This function is for configuration of nflag prior to usage.
func Configure(providedConfig ConfigOptions) {
	if providedConfig.OSSpecificFlags == false { // If we are overriding OSSpecificFlags
		Config.OSSpecificFlags = false
		Config.OSSpecificFlagString = "--" // Enforce --
	}

	if providedConfig.OSSpecificFlagString != "" { // If we are overriding OSSpecificFlagString
		Config.OSSpecificFlagString = providedConfig.OSSpecificFlagString // Set OSSpecificFlagString to whatever is provided on config
	}

	if providedConfig.ShowHelpIfNoArgs { // If ShowHelpIfNoArgs was set to true
		Config.ShowHelpIfNoArgs = true // Change to true
	}
}

// Get
// This function will get the flag value and returns it, or an error if the flag does not exist.
func Get(flagName string) (interface{}, error) {
	flag, exists := Flags[flagName] // Get the trimmedFlag and the exists bool of this flagName in Flags
	var val interface{}
	var err error

	if exists { // If the flag exists
		val = flag.Value
	} else {
		err = errors.New("Flag does not exist.")
	}

	return val, err
}

// #region GetAs functions

// GetAsBool
// This function will get the flag value and convert it to bool, or an error if the flag does not exist.
func GetAsBool(flagName string) (bool, error) {
	var flagValue bool
	flagValueInterface, flagValueError := Get(flagName)

	if flagValueError == nil { // If this flag does exist
		flagValue = flagValueInterface.(bool) // Type inference to bool
	}

	return flagValue, flagValueError
}

// GetAsInt
// This function will get the flag value and convert it to int, or an error if the flag does not exist.
func GetAsInt(flagName string) (int, error) {
	var flagValue int
	flagValueInterface, flagValueError := Get(flagName)

	if flagValueError == nil { // If this flag does exist
		flagValue = flagValueInterface.(int) // Type inference to int
	}

	return flagValue, flagValueError
}

// GetAsFloat64
// This function will get the flag value and convert it to float64, or an error if the flag does not exist.
func GetAsFloat64(flagName string) (float64, error) {
	var flagValue float64
	flagValueInterface, flagValueError := Get(flagName)

	if flagValueError == nil { // If this flag does exist
		flagValue = flagValueInterface.(float64) // Type inference to float64
	}

	return flagValue, flagValueError
}

// GetAsString
// This function will get the flag value and convert it to string, or an error if the flag does not exist.
func GetAsString(flagName string) (string, error) {
	var flagValue string
	flagValueInterface, flagValueError := Get(flagName)

	if flagValueError == nil { // If this flag does exist
		flagValue = flagValueInterface.(string) // Type inference to string
	}

	return flagValue, flagValueError
}

// #endregion

// Set
// This function is for setting a flag.
func Set(flagName string, providedFlag Flag) error {
	var errorResponse error

	if providedFlag.Type == "" { // If no Type was provided
		providedFlag.Type = "bool" // Default to bool
	}

	allowedTypes := []string{"bool", "float64", "int", "string"} // Define allowed types
	isAllowedType := false                                       // Default to providedFlag.Type not being the type

	for _, val := range allowedTypes {
		if providedFlag.Type == val { // If the Type is the same as what is allowed (this specific value)
			isAllowedType = true
		}
	}

	if isAllowedType { // If this is an allowed type
		if providedFlag.DefaultValue == nil { // If no default value was provided
			if providedFlag.Type == "bool" { // If the type is bool
				providedFlag.DefaultValue = false
			} else if providedFlag.Type == "float64" { // If the type is float64
				providedFlag.DefaultValue = 1.1
			} else if providedFlag.Type == "int" { // If the type is int
				providedFlag.DefaultValue = 1
			} else if providedFlag.Type == "string" {
				providedFlag.DefaultValue = ""
			}
		} else if (providedFlag.DefaultValue != nil) && (providedFlag.AllowNothing == false) { // If a Default value was provided
			providedFlag.AllowNothing = true                                             // Enforce AllowNothing since a default value is provided
			if providedFlag.Type != reflect.TypeOf(providedFlag.DefaultValue).String() { // If the Type and DefaultValue aren't the same type
				errorResponse = errors.New("Mismatch Flag and DefaultValue types.")
			}
		}

		if errorResponse == nil { // If there was no error
			Flags[flagName] = providedFlag // Set the flag name in Flags to the providedFlag struct
		}
	} else { // If this is a non-allowed type
		errorResponse = errors.New("Type is not allowed. Please use bool, float64, int, or string.")
	}

	return errorResponse
}

// Parse
// This function will parse input for flags
func Parse() {
	providedFlags := os.Args[1:] // ProvidedFlags start at index 1 (after binary name)

	if len(providedFlags) != 0 { // If flags were provided
		if providedFlags[0] != (Config.OSSpecificFlagString + "help") { // If the first providedFlag is not help
			for _, flag := range providedFlags { // For each flag provided
				flagNameValueSplit := strings.Split(flag, "=")                                          // Split the flag name and the value
				flagName := strings.Replace(flagNameValueSplit[0], Config.OSSpecificFlagString, "", -1) // Get the flagName and remove the OSSpecificFlagString
				flagValue := ""                                                                         // Default to nothing

				if len(flagNameValueSplit) == 2 { // If a value was provided
					flagValue = flagNameValueSplit[1] // Set to the value in flagNameValueSplit
				} else { // If the flag was passed but no value
					trimmedFlag, exists := Flags[flagName] // Get the trimmedFlag and the exists bool of this flagName in Flags

					if exists && (trimmedFlag.Type == "bool") { // If the flag exists and its type is bool
						flagValue = "true" // By passing just --flag, it implies --flag is true
					}
				}

				ParseVal(flagName, flagValue) // Parse the value
			}
		} else { // If we should output the help message
			PrintFlags() // Print the flags
			os.Exit(1)
		}

		for flagName, _ := range Flags { // For each flagName in Flags
			if Flags[flagName].Value == nil { // If no value is set for this Flag (which happens if it wasn't parsed via
				ParseVal(flagName, "flag-not-provided")
			}
		}
	} else if Config.ShowHelpIfNoArgs { // If no arguments are passed and ShowHelpIfNoArgs is true
		PrintFlags()
		os.Exit(1)
	}
}

// ParseVal
// This function will check Flags for the provided flag and parse the provided value
func ParseVal(flagName string, flagValue string) {
	trimmedFlag, exists := Flags[flagName] // Get the trimmedFlag and the exists bool of this flagName in Flags

	if exists { // If the flag exists
		if (flagValue != "") && (flagValue != "flag-not-provided") { // If a value was defined
			var conversionError error // Define conversionError as any error from not correctly converting the value to the type

			if trimmedFlag.Type == "bool" { // If the type is bool
				trimmedFlag.Value, conversionError = strconv.ParseBool(flagValue) // Convert to bool
			} else if trimmedFlag.Type == "float64" { // If the type is float64
				trimmedFlag.Value, conversionError = strconv.ParseFloat(flagValue, 64) // Convert to float64
			} else if trimmedFlag.Type == "int" { // If the type is int
				trimmedFlag.Value, conversionError = strconv.Atoi(flagValue) // Convert to int
			} else { // Remaining type, being int
				trimmedFlag.Value = flagValue
			}

			if conversionError != nil { // If there was a conversionError
				fmt.Println("An incorrect value was provided when using " + Config.OSSpecificFlagString + "flagName . Exiting.")
				os.Exit(1)
			}
		} else { // If no value was provided by the user
			returnRequiredValueMessage := false

			if flagValue == "" { // If the flag was provided by the user but no content
				if trimmedFlag.AllowNothing { // If providing no content is permitted
					if trimmedFlag.Type == "bool" { // If the flag is bool
						trimmedFlag.Value = true // Force true
					} else { // If the flag is not bool
						trimmedFlag.Value = trimmedFlag.DefaultValue // Set to defaultvalue
					}
				} else { // If providing no content is NOT permitted
					returnRequiredValueMessage = true
				}
			} else { // If this flag was not provided
				if !trimmedFlag.Required { // If the flag is NOT required
					trimmedFlag.Value = trimmedFlag.DefaultValue // Set the value to the default
				} else { // If the flag IS required
					returnRequiredValueMessage = true
				}
			}

			if returnRequiredValueMessage {
				fmt.Println("A required value for " + Config.OSSpecificFlagString + flagName + " was not provided. Exiting.")
				os.Exit(1)
			}
		}

		Flags[flagName] = trimmedFlag // Update the map
	} else { // If the flag does not exist.
		fmt.Println(Config.OSSpecificFlagString + flagName + " does not exist. Exiting.")
		os.Exit(1)
	}
}

// PrintFlags
// This function will print all the flags that are set and their defaults
func PrintFlags() {
	fmt.Println("Usage: " + Config.OSSpecificFlagString + "example=value")
	fmt.Println("The following options are available:")

	for flagName, flag := range Flags { // For each flagName and trimmedFlag struct in Flags
		fmt.Println(Config.OSSpecificFlagString + flagName + " " + flag.Type)
		fmt.Println("\t" + flag.Descriptor)

		if flag.DefaultValue != nil { // If DefaultValue is not nil
			fmt.Println("\t" + "Default Value: " + reflect.ValueOf(flag.DefaultValue).String())
		}

		fmt.Println("\t" + "Allows Providing Only Flag: " + strconv.FormatBool(flag.AllowNothing))
		fmt.Println("")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/codegangsta/cli"
)

func execProgram(program string, input string) string {
	var stdout bytes.Buffer
	cmd := exec.Command(program)
	cmd.Stdin = strings.NewReader(input)
	cmd.Stdout = &stdout

	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	return stdout.String()
}

func main() {
	app := cli.NewApp()
	app.Name = "nigit"
	app.Version = "0.1-alpha"
	app.Usage = "Expose any Program as HTTP API"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "port",
			Value:  "8000",
			Usage:  "HTTP port",
			EnvVar: "PORT",
		},
	}

	app.Action = func(c *cli.Context) {
		program := c.Args().First()
		if program == "" {
			fmt.Println("Please provide name of the script to run under nigit")
			os.Exit(1)
		}

		program, err := filepath.Abs(program)
		if err != nil {
			fmt.Printf("Cannot get path of %s\n", program)
			os.Exit(2)
		}

		programPath, err := exec.LookPath(program)
		if err != nil {
			fmt.Printf("Executable program %s not found\n", program)
			os.Exit(3)
		}

		serve := func(w http.ResponseWriter, r *http.Request) {
			buf := new(bytes.Buffer)
			buf.ReadFrom(r.Body)

			stdout := execProgram(programPath, buf.String())
			io.WriteString(w, stdout)
		}

		http.HandleFunc("/", serve)
		http.ListenAndServe(":"+c.GlobalString("port"), nil)
	}

	app.Run(os.Args)
}
<commit_msg>Support multiple files<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/codegangsta/cli"
)

func execProgram(program string, input string) string {
	var stdout bytes.Buffer
	cmd := exec.Command(program)
	cmd.Stdin = strings.NewReader(input)
	cmd.Stdout = &stdout

	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}

	return stdout.String()
}

func urlPath(programPath string) string {
	return "/" + strings.TrimSuffix(filepath.Base(programPath), filepath.Ext(programPath))
}

func main() {
	app := cli.NewApp()
	app.Name = "nigit"
	app.Version = "0.1-alpha"
	app.Usage = "Expose any Program as HTTP API"
	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "port",
			Value:  "8000",
			Usage:  "HTTP port",
			EnvVar: "PORT",
		},
	}

	app.Action = func(c *cli.Context) {
		if !c.Args().Present() {
			fmt.Println("Please provide the names of the scripts to run under nigit")
			os.Exit(1)
		}
		fmt.Println("Serve from port " + c.GlobalString("port"))

		for _, program := range c.Args() {
			programPath, err := filepath.Abs(program)
			if err != nil {
				fmt.Printf("Cannot get path of %s\n", program)
				os.Exit(2)
			}

			programPath, err = exec.LookPath(programPath)
			if err != nil {
				fmt.Printf("Executable program %s not found\n", program)
				os.Exit(3)
			}

			serve := func(w http.ResponseWriter, r *http.Request) {
				buf := new(bytes.Buffer)
				buf.ReadFrom(r.Body)

				stdout := execProgram(programPath, buf.String())
				io.WriteString(w, stdout)
			}

			fmt.Printf("%s -> %s\n", urlPath(programPath), program)
			http.HandleFunc(urlPath(programPath), serve)
		}
		http.ListenAndServe(":"+c.GlobalString("port"), nil)
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package gocassa

import (
	"fmt"
	"strings"
	"time"
)

type tableFactory interface {
	NewTable(string, interface{}, map[string]interface{}, Keys) Table
}

type k struct {
	qe           QueryExecutor
	name         string
	debugMode    bool
	tableFactory tableFactory
}

// Connect to a certain keyspace directly. Same as using Connect().KeySpace(keySpaceName)
func ConnectToKeySpace(keySpace string, nodeIps []string, username, password string) (KeySpace, error) {
	c, err := Connect(nodeIps, username, password)
	if err != nil {
		return nil, err
	}
	return c.KeySpace(keySpace), nil
}

func (k *k) DebugMode(b bool) {
	k.debugMode = b
}

func (k *k) Table(name string, entity interface{}, keys Keys) Table {
	n := name + "__" + strings.Join(keys.PartitionKeys, "_") + "__" + strings.Join(keys.ClusteringColumns, "_")
	m, ok := toMap(entity)
	if !ok {
		panic("Unrecognized row type")
	}
	return k.NewTable(n, entity, m, keys)
}

func (k *k) NewTable(name string, entity interface{}, fields map[string]interface{}, keys Keys) Table {
	// Act both as a proxy to a tableFactory, and as the tableFactory itself (in most situations, a k will be its own
	// tableFactory, but not always [ie. mocking])
	if k.tableFactory != k {
		return k.tableFactory.NewTable(name, entity, fields, keys)
	} else {
		ti := newTableInfo(k.name, name, keys, entity, fields)
		return &t{
			keySpace: k,
			info:     ti,
			options:  Options{},
		}
	}
}

func (k *k) MapTable(name, id string, row interface{}) MapTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	return &mapT{
		t: k.NewTable(fmt.Sprintf("%s_map_%s", name, id), row, m, Keys{
			PartitionKeys: []string{id},
		}),
		idField: id,
	}
}

func (k *k) SetKeysSpaceName(name string) {
	k.name = name
}

func (k *k) MultimapTable(name, fieldToIndexBy, id string, row interface{}) MultimapTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	return &multimapT{
		t: k.NewTable(fmt.Sprintf("%s_multimap_%s_%s", name, fieldToIndexBy, id), row, m, Keys{
			PartitionKeys:     []string{fieldToIndexBy},
			ClusteringColumns: []string{id},
		}),
		idField:        id,
		fieldToIndexBy: fieldToIndexBy,
	}
}

func (k *k) MultimapMultiKeyTable(name string, fieldToIndexBy, id []string, row interface{}) MultimapMkTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	return &multimapMkT{
		t: k.NewTable(fmt.Sprintf("%s_multimapMk", name), row, m, Keys{
			PartitionKeys:     fieldToIndexBy,
			ClusteringColumns: id,
		}),
		idField:         id,
		fieldsToIndexBy: fieldToIndexBy,
	}
}

func (k *k) TimeSeriesTable(name, timeField, idField string, bucketSize time.Duration, row interface{}) TimeSeriesTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	m[bucketFieldName] = time.Now()
	return &timeSeriesT{
		t: k.NewTable(fmt.Sprintf("%s_timeSeries_%s_%s_%s", name, timeField, idField, bucketSize), row, m, Keys{
			PartitionKeys:     []string{bucketFieldName},
			ClusteringColumns: []string{timeField, idField},
		}),
		timeField:  timeField,
		idField:    idField,
		bucketSize: bucketSize,
	}
}

func (k *k) MultiTimeSeriesTable(name, indexField, timeField, idField string, bucketSize time.Duration, row interface{}) MultiTimeSeriesTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	m[bucketFieldName] = time.Now()
	return &multiTimeSeriesT{
		t: k.NewTable(fmt.Sprintf("%s_multiTimeSeries_%s_%s_%s_%s", name, indexField, timeField, idField, bucketSize.String()), row, m, Keys{
			PartitionKeys:     []string{indexField, bucketFieldName},
			ClusteringColumns: []string{timeField, idField},
		}),
		indexField: indexField,
		timeField:  timeField,
		idField:    idField,
		bucketSize: bucketSize,
	}
}

func (k *k) FlakeSeriesTable(name, idField string, bucketSize time.Duration, row interface{}) FlakeSeriesTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	m[flakeTimestampFieldName] = time.Now()
	m[bucketFieldName] = time.Now()
	return &flakeSeriesT{
		t: k.NewTable(fmt.Sprintf("%s_flakeSeries_%s_%s", name, idField, bucketSize.String()), row, m, Keys{
			PartitionKeys:     []string{bucketFieldName},
			ClusteringColumns: []string{flakeTimestampFieldName, idField},
		}),
		idField:    idField,
		bucketSize: bucketSize,
	}
}

func (k *k) MultiFlakeSeriesTable(name, indexField, idField string, bucketSize time.Duration, row interface{}) MultiFlakeSeriesTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	m[flakeTimestampFieldName] = time.Now()
	m[bucketFieldName] = time.Now()
	return &multiFlakeSeriesT{
		t: k.NewTable(fmt.Sprintf("%s_multiflakeSeries_%s_%s_%s", name, indexField, idField, bucketSize.String()), row, m, Keys{
			PartitionKeys:     []string{indexField, bucketFieldName},
			ClusteringColumns: []string{flakeTimestampFieldName, idField},
		}),
		idField:    idField,
		bucketSize: bucketSize,
		indexField: indexField,
	}
}

// Returns table names in a keyspace
func (k *k) Tables() ([]string, error) {
	const stmt = "SELECT columnfamily_name FROM system.schema_columnfamilies WHERE keyspace_name = ?"
	maps, err := k.qe.Query(stmt, k.name)
	if err != nil {
		return nil, err
	}
	ret := []string{}
	for _, m := range maps {
		ret = append(ret, m["columnfamily_name"].(string))
	}
	return ret, nil
}

func (k *k) Exists(cf string) (bool, error) {
	ts, err := k.Tables()
	if err != nil {
		return false, err
	}
	for _, v := range ts {
		if strings.ToLower(v) == strings.ToLower(cf) {
			return true, nil
		}
	}
	return false, nil
}

func (k *k) DropTable(cf string) error {
	stmt := fmt.Sprintf("DROP TABLE IF EXISTS %s.%s", k.name, cf)
	return k.qe.Execute(stmt)
}

func (k *k) Name() string {
	return k.name
}
<commit_msg>Return error if no query executor is available<commit_after>package gocassa

import (
	"fmt"
	"strings"
	"time"
)

type tableFactory interface {
	NewTable(string, interface{}, map[string]interface{}, Keys) Table
}

type k struct {
	qe           QueryExecutor
	name         string
	debugMode    bool
	tableFactory tableFactory
}

// Connect to a certain keyspace directly. Same as using Connect().KeySpace(keySpaceName)
func ConnectToKeySpace(keySpace string, nodeIps []string, username, password string) (KeySpace, error) {
	c, err := Connect(nodeIps, username, password)
	if err != nil {
		return nil, err
	}
	return c.KeySpace(keySpace), nil
}

func (k *k) DebugMode(b bool) {
	k.debugMode = b
}

func (k *k) Table(name string, entity interface{}, keys Keys) Table {
	n := name + "__" + strings.Join(keys.PartitionKeys, "_") + "__" + strings.Join(keys.ClusteringColumns, "_")
	m, ok := toMap(entity)
	if !ok {
		panic("Unrecognized row type")
	}
	return k.NewTable(n, entity, m, keys)
}

func (k *k) NewTable(name string, entity interface{}, fields map[string]interface{}, keys Keys) Table {
	// Act both as a proxy to a tableFactory, and as the tableFactory itself (in most situations, a k will be its own
	// tableFactory, but not always [ie. mocking])
	if k.tableFactory != k {
		return k.tableFactory.NewTable(name, entity, fields, keys)
	} else {
		ti := newTableInfo(k.name, name, keys, entity, fields)
		return &t{
			keySpace: k,
			info:     ti,
			options:  Options{},
		}
	}
}

func (k *k) MapTable(name, id string, row interface{}) MapTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	return &mapT{
		t: k.NewTable(fmt.Sprintf("%s_map_%s", name, id), row, m, Keys{
			PartitionKeys: []string{id},
		}),
		idField: id,
	}
}

func (k *k) SetKeysSpaceName(name string) {
	k.name = name
}

func (k *k) MultimapTable(name, fieldToIndexBy, id string, row interface{}) MultimapTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	return &multimapT{
		t: k.NewTable(fmt.Sprintf("%s_multimap_%s_%s", name, fieldToIndexBy, id), row, m, Keys{
			PartitionKeys:     []string{fieldToIndexBy},
			ClusteringColumns: []string{id},
		}),
		idField:        id,
		fieldToIndexBy: fieldToIndexBy,
	}
}

func (k *k) MultimapMultiKeyTable(name string, fieldToIndexBy, id []string, row interface{}) MultimapMkTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	return &multimapMkT{
		t: k.NewTable(fmt.Sprintf("%s_multimapMk", name), row, m, Keys{
			PartitionKeys:     fieldToIndexBy,
			ClusteringColumns: id,
		}),
		idField:         id,
		fieldsToIndexBy: fieldToIndexBy,
	}
}

func (k *k) TimeSeriesTable(name, timeField, idField string, bucketSize time.Duration, row interface{}) TimeSeriesTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	m[bucketFieldName] = time.Now()
	return &timeSeriesT{
		t: k.NewTable(fmt.Sprintf("%s_timeSeries_%s_%s_%s", name, timeField, idField, bucketSize), row, m, Keys{
			PartitionKeys:     []string{bucketFieldName},
			ClusteringColumns: []string{timeField, idField},
		}),
		timeField:  timeField,
		idField:    idField,
		bucketSize: bucketSize,
	}
}

func (k *k) MultiTimeSeriesTable(name, indexField, timeField, idField string, bucketSize time.Duration, row interface{}) MultiTimeSeriesTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	m[bucketFieldName] = time.Now()
	return &multiTimeSeriesT{
		t: k.NewTable(fmt.Sprintf("%s_multiTimeSeries_%s_%s_%s_%s", name, indexField, timeField, idField, bucketSize.String()), row, m, Keys{
			PartitionKeys:     []string{indexField, bucketFieldName},
			ClusteringColumns: []string{timeField, idField},
		}),
		indexField: indexField,
		timeField:  timeField,
		idField:    idField,
		bucketSize: bucketSize,
	}
}

func (k *k) FlakeSeriesTable(name, idField string, bucketSize time.Duration, row interface{}) FlakeSeriesTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	m[flakeTimestampFieldName] = time.Now()
	m[bucketFieldName] = time.Now()
	return &flakeSeriesT{
		t: k.NewTable(fmt.Sprintf("%s_flakeSeries_%s_%s", name, idField, bucketSize.String()), row, m, Keys{
			PartitionKeys:     []string{bucketFieldName},
			ClusteringColumns: []string{flakeTimestampFieldName, idField},
		}),
		idField:    idField,
		bucketSize: bucketSize,
	}
}

func (k *k) MultiFlakeSeriesTable(name, indexField, idField string, bucketSize time.Duration, row interface{}) MultiFlakeSeriesTable {
	m, ok := toMap(row)
	if !ok {
		panic("Unrecognized row type")
	}
	m[flakeTimestampFieldName] = time.Now()
	m[bucketFieldName] = time.Now()
	return &multiFlakeSeriesT{
		t: k.NewTable(fmt.Sprintf("%s_multiflakeSeries_%s_%s_%s", name, indexField, idField, bucketSize.String()), row, m, Keys{
			PartitionKeys:     []string{indexField, bucketFieldName},
			ClusteringColumns: []string{flakeTimestampFieldName, idField},
		}),
		idField:    idField,
		bucketSize: bucketSize,
		indexField: indexField,
	}
}

// Returns table names in a keyspace
func (k *k) Tables() ([]string, error) {
	const stmt = "SELECT columnfamily_name FROM system.schema_columnfamilies WHERE keyspace_name = ?"

	if k.qe == nil {
		return nil, fmt.Errorf("no query executor configured")
	}

	maps, err := k.qe.Query(stmt, k.name)
	if err != nil {
		return nil, err
	}
	ret := []string{}
	for _, m := range maps {
		ret = append(ret, m["columnfamily_name"].(string))
	}
	return ret, nil
}

func (k *k) Exists(cf string) (bool, error) {
	ts, err := k.Tables()
	if err != nil {
		return false, err
	}
	for _, v := range ts {
		if strings.ToLower(v) == strings.ToLower(cf) {
			return true, nil
		}
	}
	return false, nil
}

func (k *k) DropTable(cf string) error {
	stmt := fmt.Sprintf("DROP TABLE IF EXISTS %s.%s", k.name, cf)
	return k.qe.Execute(stmt)
}

func (k *k) Name() string {
	return k.name
}
<|endoftext|>


<commit_before>package http

import (
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"v2ray.com/core/app"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/bufio"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/log"
	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/common/serial"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/ray"
)

// Server is a HTTP proxy server.
type Server struct {
	sync.Mutex
	accepting        bool
	packetDispatcher dispatcher.PacketDispatcher
	config           *ServerConfig
	tcpListener      *internet.TCPHub
	meta             *proxy.InboundHandlerMeta
}

func NewServer(config *ServerConfig, packetDispatcher dispatcher.PacketDispatcher, meta *proxy.InboundHandlerMeta) *Server {
	return &Server{
		packetDispatcher: packetDispatcher,
		config:           config,
		meta:             meta,
	}
}

func (v *Server) Port() v2net.Port {
	return v.meta.Port
}

func (v *Server) Close() {
	v.accepting = false
	if v.tcpListener != nil {
		v.Lock()
		v.tcpListener.Close()
		v.tcpListener = nil
		v.Unlock()
	}
}

func (v *Server) Start() error {
	if v.accepting {
		return nil
	}

	tcpListener, err := internet.ListenTCP(v.meta.Address, v.meta.Port, v.handleConnection, v.meta.StreamSettings)
	if err != nil {
		log.Error("HTTP: Failed listen on ", v.meta.Address, ":", v.meta.Port, ": ", err)
		return err
	}
	v.Lock()
	v.tcpListener = tcpListener
	v.Unlock()
	v.accepting = true
	return nil
}

func parseHost(rawHost string, defaultPort v2net.Port) (v2net.Destination, error) {
	port := defaultPort
	host, rawPort, err := net.SplitHostPort(rawHost)
	if err != nil {
		if addrError, ok := err.(*net.AddrError); ok && strings.Contains(addrError.Err, "missing port") {
			host = rawHost
		} else {
			return v2net.Destination{}, err
		}
	} else {
		intPort, err := strconv.Atoi(rawPort)
		if err != nil {
			return v2net.Destination{}, err
		}
		port = v2net.Port(intPort)
	}

	if ip := net.ParseIP(host); ip != nil {
		return v2net.TCPDestination(v2net.IPAddress(ip), port), nil
	}
	return v2net.TCPDestination(v2net.DomainAddress(host), port), nil
}

func (v *Server) handleConnection(conn internet.Connection) {
	defer conn.Close()
	timedReader := v2net.NewTimeOutReader(v.config.Timeout, conn)
	reader := bufio.OriginalReaderSize(timedReader, 2048)

	request, err := http.ReadRequest(reader)
	if err != nil {
		if errors.Cause(err) != io.EOF {
			log.Warning("HTTP: Failed to read http request: ", err)
		}
		return
	}
	log.Info("HTTP: Request to Method [", request.Method, "] Host [", request.Host, "] with URL [", request.URL, "]")
	defaultPort := v2net.Port(80)
	if strings.ToLower(request.URL.Scheme) == "https" {
		defaultPort = v2net.Port(443)
	}
	host := request.Host
	if len(host) == 0 {
		host = request.URL.Host
	}
	dest, err := parseHost(host, defaultPort)
	if err != nil {
		log.Warning("HTTP: Malformed proxy host (", host, "): ", err)
		return
	}
	log.Access(conn.RemoteAddr(), request.URL, log.AccessAccepted, "")
	session := &proxy.SessionInfo{
		Source:      v2net.DestinationFromAddr(conn.RemoteAddr()),
		Destination: dest,
		Inbound:     v.meta,
	}
	if strings.ToUpper(request.Method) == "CONNECT" {
		v.handleConnect(request, session, reader, conn)
	} else {
		v.handlePlainHTTP(request, session, reader, conn)
	}
}

func (v *Server) handleConnect(request *http.Request, session *proxy.SessionInfo, reader io.Reader, writer io.Writer) {
	response := &http.Response{
		Status:        "200 OK",
		StatusCode:    200,
		Proto:         "HTTP/1.1",
		ProtoMajor:    1,
		ProtoMinor:    1,
		Header:        http.Header(make(map[string][]string)),
		Body:          nil,
		ContentLength: 0,
		Close:         false,
	}
	response.Write(writer)

	ray := v.packetDispatcher.DispatchToOutbound(session)
	v.transport(reader, writer, ray)
}

func (v *Server) transport(input io.Reader, output io.Writer, ray ray.InboundRay) {
	var wg sync.WaitGroup
	wg.Add(2)
	defer wg.Wait()

	go func() {
		v2reader := buf.NewReader(input)
		defer v2reader.Release()

		if err := buf.PipeUntilEOF(v2reader, ray.InboundInput()); err != nil {
			log.Info("HTTP: Failed to transport all TCP request: ", err)
		}
		ray.InboundInput().Close()
		wg.Done()
	}()

	go func() {
		v2writer := buf.NewWriter(output)
		defer v2writer.Release()

		if err := buf.PipeUntilEOF(ray.InboundOutput(), v2writer); err != nil {
			log.Info("HTTP: Failed to transport all TCP response: ", err)
		}
		ray.InboundOutput().Release()
		wg.Done()
	}()
}

// @VisibleForTesting
func StripHopByHopHeaders(request *http.Request) {
	// Strip hop-by-hop header basaed on RFC:
	// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html#sec13.5.1
	// https://www.mnot.net/blog/2011/07/11/what_proxies_must_do

	request.Header.Del("Proxy-Connection")
	request.Header.Del("Proxy-Authenticate")
	request.Header.Del("Proxy-Authorization")
	request.Header.Del("TE")
	request.Header.Del("Trailers")
	request.Header.Del("Transfer-Encoding")
	request.Header.Del("Upgrade")

	// TODO: support keep-alive
	connections := request.Header.Get("Connection")
	request.Header.Set("Connection", "close")
	if len(connections) == 0 {
		return
	}
	for _, h := range strings.Split(connections, ",") {
		request.Header.Del(strings.TrimSpace(h))
	}
}

func (v *Server) GenerateResponse(statusCode int, status string) *http.Response {
	hdr := http.Header(make(map[string][]string))
	hdr.Set("Connection", "close")
	return &http.Response{
		Status:        status,
		StatusCode:    statusCode,
		Proto:         "HTTP/1.1",
		ProtoMajor:    1,
		ProtoMinor:    1,
		Header:        hdr,
		Body:          nil,
		ContentLength: 0,
		Close:         false,
	}
}

func (v *Server) handlePlainHTTP(request *http.Request, session *proxy.SessionInfo, reader io.Reader, writer io.Writer) {
	if len(request.URL.Host) <= 0 {
		response := v.GenerateResponse(400, "Bad Request")
		response.Write(writer)

		return
	}

	request.Host = request.URL.Host
	StripHopByHopHeaders(request)

	ray := v.packetDispatcher.DispatchToOutbound(session)
	defer ray.InboundInput().Close()
	defer ray.InboundOutput().Release()

	var finish sync.WaitGroup
	finish.Add(1)
	go func() {
		defer finish.Done()
		requestWriter := bufio.NewWriter(buf.NewBytesWriter(ray.InboundInput()))
		err := request.Write(requestWriter)
		if err != nil {
			log.Warning("HTTP: Failed to write request: ", err)
			return
		}
		requestWriter.Flush()
	}()

	finish.Add(1)
	go func() {
		defer finish.Done()
		responseReader := bufio.OriginalReader(buf.NewBytesReader(ray.InboundOutput()))
		response, err := http.ReadResponse(responseReader, request)
		if err != nil {
			log.Warning("HTTP: Failed to read response: ", err)
			response = v.GenerateResponse(503, "Service Unavailable")
		}
		responseWriter := bufio.NewWriter(writer)
		err = response.Write(responseWriter)
		if err != nil {
			log.Warning("HTTP: Failed to write response: ", err)
			return
		}
		responseWriter.Flush()
	}()
	finish.Wait()
}

type ServerFactory struct{}

func (v *ServerFactory) StreamCapability() v2net.NetworkList {
	return v2net.NetworkList{
		Network: []v2net.Network{v2net.Network_RawTCP},
	}
}

func (v *ServerFactory) Create(space app.Space, rawConfig interface{}, meta *proxy.InboundHandlerMeta) (proxy.InboundHandler, error) {
	if !space.HasApp(dispatcher.APP_ID) {
		return nil, common.ErrBadConfiguration
	}
	return NewServer(
		rawConfig.(*ServerConfig),
		space.GetApp(dispatcher.APP_ID).(dispatcher.PacketDispatcher),
		meta), nil
}

func init() {
	proxy.MustRegisterInboundHandlerCreator(serial.GetMessageType(new(ServerConfig)), new(ServerFactory))
}
<commit_msg>Close http connection ASAP.<commit_after>package http

import (
	"io"
	"net"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"v2ray.com/core/app"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/bufio"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/log"
	v2net "v2ray.com/core/common/net"
	"v2ray.com/core/common/serial"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/ray"
)

// Server is a HTTP proxy server.
type Server struct {
	sync.Mutex
	accepting        bool
	packetDispatcher dispatcher.PacketDispatcher
	config           *ServerConfig
	tcpListener      *internet.TCPHub
	meta             *proxy.InboundHandlerMeta
}

// NewServer creates a new HTTP inbound handler.
func NewServer(config *ServerConfig, packetDispatcher dispatcher.PacketDispatcher, meta *proxy.InboundHandlerMeta) *Server {
	return &Server{
		packetDispatcher: packetDispatcher,
		config:           config,
		meta:             meta,
	}
}

// Port implements InboundHandler.Port().
func (v *Server) Port() v2net.Port {
	return v.meta.Port
}

// Close implements InboundHandler.Close().
func (v *Server) Close() {
	v.accepting = false
	if v.tcpListener != nil {
		v.Lock()
		v.tcpListener.Close()
		v.tcpListener = nil
		v.Unlock()
	}
}

// Start implements InboundHandler.Start().
func (v *Server) Start() error {
	if v.accepting {
		return nil
	}

	tcpListener, err := internet.ListenTCP(v.meta.Address, v.meta.Port, v.handleConnection, v.meta.StreamSettings)
	if err != nil {
		log.Error("HTTP: Failed listen on ", v.meta.Address, ":", v.meta.Port, ": ", err)
		return err
	}
	v.Lock()
	v.tcpListener = tcpListener
	v.Unlock()
	v.accepting = true
	return nil
}

func parseHost(rawHost string, defaultPort v2net.Port) (v2net.Destination, error) {
	port := defaultPort
	host, rawPort, err := net.SplitHostPort(rawHost)
	if err != nil {
		if addrError, ok := err.(*net.AddrError); ok && strings.Contains(addrError.Err, "missing port") {
			host = rawHost
		} else {
			return v2net.Destination{}, err
		}
	} else {
		intPort, err := strconv.Atoi(rawPort)
		if err != nil {
			return v2net.Destination{}, err
		}
		port = v2net.Port(intPort)
	}

	if ip := net.ParseIP(host); ip != nil {
		return v2net.TCPDestination(v2net.IPAddress(ip), port), nil
	}
	return v2net.TCPDestination(v2net.DomainAddress(host), port), nil
}

func (v *Server) handleConnection(conn internet.Connection) {
	defer conn.Close()
	timedReader := v2net.NewTimeOutReader(v.config.Timeout, conn)
	reader := bufio.OriginalReaderSize(timedReader, 2048)

	request, err := http.ReadRequest(reader)
	if err != nil {
		if errors.Cause(err) != io.EOF {
			log.Warning("HTTP: Failed to read http request: ", err)
		}
		return
	}
	log.Info("HTTP: Request to Method [", request.Method, "] Host [", request.Host, "] with URL [", request.URL, "]")
	defaultPort := v2net.Port(80)
	if strings.ToLower(request.URL.Scheme) == "https" {
		defaultPort = v2net.Port(443)
	}
	host := request.Host
	if len(host) == 0 {
		host = request.URL.Host
	}
	dest, err := parseHost(host, defaultPort)
	if err != nil {
		log.Warning("HTTP: Malformed proxy host (", host, "): ", err)
		return
	}
	log.Access(conn.RemoteAddr(), request.URL, log.AccessAccepted, "")
	session := &proxy.SessionInfo{
		Source:      v2net.DestinationFromAddr(conn.RemoteAddr()),
		Destination: dest,
		Inbound:     v.meta,
	}
	if strings.ToUpper(request.Method) == "CONNECT" {
		v.handleConnect(request, session, reader, conn)
	} else {
		v.handlePlainHTTP(request, session, reader, conn)
	}
}

func (v *Server) handleConnect(request *http.Request, session *proxy.SessionInfo, reader io.Reader, writer io.Writer) {
	response := &http.Response{
		Status:        "200 OK",
		StatusCode:    200,
		Proto:         "HTTP/1.1",
		ProtoMajor:    1,
		ProtoMinor:    1,
		Header:        http.Header(make(map[string][]string)),
		Body:          nil,
		ContentLength: 0,
		Close:         false,
	}
	if err := response.Write(writer); err != nil {
		log.Warning("HTTP|Server: failed to write back OK response: ", err)
		return
	}

	ray := v.packetDispatcher.DispatchToOutbound(session)
	v.transport(reader, writer, ray)
}

func (v *Server) transport(input io.Reader, output io.Writer, ray ray.InboundRay) {
	var wg sync.WaitGroup
	wg.Add(2)
	defer wg.Wait()

	go func() {
		v2reader := buf.NewReader(input)
		defer v2reader.Release()

		if err := buf.PipeUntilEOF(v2reader, ray.InboundInput()); err != nil {
			log.Info("HTTP: Failed to transport all TCP request: ", err)
		}
		ray.InboundInput().Close()
		wg.Done()
	}()

	go func() {
		v2writer := buf.NewWriter(output)
		defer v2writer.Release()

		if err := buf.PipeUntilEOF(ray.InboundOutput(), v2writer); err != nil {
			log.Info("HTTP: Failed to transport all TCP response: ", err)
		}
		ray.InboundOutput().Release()
		wg.Done()
	}()
}

// @VisibleForTesting
func StripHopByHopHeaders(request *http.Request) {
	// Strip hop-by-hop header basaed on RFC:
	// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html#sec13.5.1
	// https://www.mnot.net/blog/2011/07/11/what_proxies_must_do

	request.Header.Del("Proxy-Connection")
	request.Header.Del("Proxy-Authenticate")
	request.Header.Del("Proxy-Authorization")
	request.Header.Del("TE")
	request.Header.Del("Trailers")
	request.Header.Del("Transfer-Encoding")
	request.Header.Del("Upgrade")

	// TODO: support keep-alive
	connections := request.Header.Get("Connection")
	request.Header.Set("Connection", "close")
	if len(connections) == 0 {
		return
	}
	for _, h := range strings.Split(connections, ",") {
		request.Header.Del(strings.TrimSpace(h))
	}
}

func (v *Server) GenerateResponse(statusCode int, status string) *http.Response {
	hdr := http.Header(make(map[string][]string))
	hdr.Set("Connection", "close")
	return &http.Response{
		Status:        status,
		StatusCode:    statusCode,
		Proto:         "HTTP/1.1",
		ProtoMajor:    1,
		ProtoMinor:    1,
		Header:        hdr,
		Body:          nil,
		ContentLength: 0,
		Close:         false,
	}
}

func (v *Server) handlePlainHTTP(request *http.Request, session *proxy.SessionInfo, reader io.Reader, writer io.Writer) {
	if len(request.URL.Host) <= 0 {
		response := v.GenerateResponse(400, "Bad Request")
		response.Write(writer)

		return
	}

	request.Host = request.URL.Host
	StripHopByHopHeaders(request)

	ray := v.packetDispatcher.DispatchToOutbound(session)

	var finish sync.WaitGroup
	finish.Add(1)
	go func() {
		defer ray.InboundInput().Close()
		defer finish.Done()
		requestWriter := bufio.NewWriter(buf.NewBytesWriter(ray.InboundInput()))
		err := request.Write(requestWriter)
		if err != nil {
			log.Warning("HTTP: Failed to write request: ", err)
			return
		}
		requestWriter.Flush()
	}()

	finish.Add(1)
	go func() {
		defer ray.InboundOutput().Release()
		defer finish.Done()
		responseReader := bufio.OriginalReader(buf.NewBytesReader(ray.InboundOutput()))
		response, err := http.ReadResponse(responseReader, request)
		if err != nil {
			log.Warning("HTTP: Failed to read response: ", err)
			response = v.GenerateResponse(503, "Service Unavailable")
		}
		responseWriter := bufio.NewWriter(writer)
		err = response.Write(responseWriter)
		if err != nil {
			log.Warning("HTTP: Failed to write response: ", err)
			return
		}
		responseWriter.Flush()
	}()

	finish.Wait()
}

// ServerFactory is a InboundHandlerFactory.
type ServerFactory struct{}

// StreamCapability implements InboundHandlerFactory.StreamCapability().
func (v *ServerFactory) StreamCapability() v2net.NetworkList {
	return v2net.NetworkList{
		Network: []v2net.Network{v2net.Network_RawTCP},
	}
}

// Create implements InboundHandlerFactory.Create().
func (v *ServerFactory) Create(space app.Space, rawConfig interface{}, meta *proxy.InboundHandlerMeta) (proxy.InboundHandler, error) {
	if !space.HasApp(dispatcher.APP_ID) {
		return nil, common.ErrBadConfiguration
	}
	return NewServer(
		rawConfig.(*ServerConfig),
		space.GetApp(dispatcher.APP_ID).(dispatcher.PacketDispatcher),
		meta), nil
}

func init() {
	proxy.MustRegisterInboundHandlerCreator(serial.GetMessageType(new(ServerConfig)), new(ServerFactory))
}
<|endoftext|>


<commit_before>package v2ray

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"github.com/miekg/dns"
	vcore "v2ray.com/core"
	vnet "v2ray.com/core/common/net"
	vsession "v2ray.com/core/common/session"
	vdns "v2ray.com/core/features/dns"

	"github.com/eycorsican/go-tun2socks/core"
	"github.com/eycorsican/go-tun2socks/proxy"
)

func isIPv4(ip net.IP) bool {
	if ip.To4() != nil {
		return true
	}
	return false
}

func isIPv6(ip net.IP) bool {
	// To16() also valid for ipv4, ensure it's not an ipv4 address
	if ip.To4() != nil {
		return false
	}
	if ip.To16() != nil {
		return true
	}
	return false
}

type connEntry struct {
	conn                net.Conn
	target              vnet.Destination
	cancelFetchingInput context.CancelFunc
	fetchingInputCtx    context.Context
}

type dnsRespEntry struct {
	err  error
	data []byte
	conn core.Connection
}

type handler struct {
	sync.Mutex

	ctx        context.Context
	v          *vcore.Instance
	conns      map[core.Connection]*connEntry
	dispatched map[core.Connection]bool
	dnsRespCh  chan *dnsRespEntry
	dnsClient  vdns.Client
}

func (h *handler) shouldAcceptDNSQuery(data []byte) bool {
	req := new(dns.Msg)
	err := req.Unpack(data)
	if err != nil {
		return false
	}

	// TODO: allow multiple question
	if len(req.Question) != 1 {
		return false
	}

	qtype := req.Question[0].Qtype
	if qtype != dns.TypeA && qtype != dns.TypeAAAA {
		return false
	}

	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	if _, ok := dns.IsDomainName(domain); !ok {
		return false
	}

	return true
}

func (h *handler) handleDNSQuery(conn core.Connection, data []byte) {
	var err error
	var answer []byte = nil
	defer func() {
		h.dnsRespCh <- &dnsRespEntry{conn: conn, data: answer, err: err}
	}()

	// No error checks here because they are already done in shouldAcceptDNSQuery()
	req := new(dns.Msg)
	req.Unpack(data)
	qtype := req.Question[0].Qtype
	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	log.Printf("dispatch dns request for domain: %v (%v)", domain, qtype)
	ips, err := h.dnsClient.LookupIP(domain)
	if err != nil {
		err = errors.New(fmt.Sprintf("lookup ip failed: %v", err))
		return
	}

	resp := new(dns.Msg)
	resp.SetReply(req)
	resp.RecursionAvailable = true
	for _, ip := range ips {
		if isIPv4(ip) && qtype == dns.TypeA {
			resp.Answer = append(resp.Answer, &dns.A{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 4,
				},
				A: ip,
			})
		} else if isIPv6(ip) && qtype == dns.TypeAAAA {
			resp.Answer = append(resp.Answer, &dns.AAAA{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeAAAA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 16,
				},
				AAAA: ip,
			})
		}
	}
	if len(resp.Answer) == 0 {
		err = errors.New("no answer")
		return
	}
	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)
	dnsAnswer, err := resp.PackBuffer(buf)
	if err != nil {
		err = errors.New(fmt.Sprintf("packing dns resp msg failed: %v", err))
	} else {
		answer = append([]byte(nil), dnsAnswer...)
	}
	return
}

func (h *handler) handleDNSResponse() {
	for {
		select {
		case respEntry := <-h.dnsRespCh:
			if respEntry.err == nil {
				_, err := respEntry.conn.Write(respEntry.data)
				if err != nil {
					log.Printf("write dns response to local failed: %v", err)
				}
			} else {
				log.Printf("dispatch dns request failed: %v", respEntry.err)
			}
			h.Close(respEntry.conn)
			respEntry.conn.Close()
		}
	}
}

func (h *handler) fetchInput(conn core.Connection) {
	h.Lock()
	c, ok := h.conns[conn]
	h.Unlock()
	if !ok {
		return
	}

	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)

FetchingLoop:
	for {
		c.conn.SetReadDeadline(time.Now().Add(4 * time.Second))
		n, err := c.conn.Read(buf)
		if err, ok := err.(net.Error); ok && err.Timeout() {
			select {
			case <-c.fetchingInputCtx.Done():
				// Request was handed to V2Ray, stop fetching but leave the
				// connection open.
				return
			default:
				continue FetchingLoop
			}
		}
		if err != nil {
			h.Close(conn)
			conn.Close()
			return
		}
		_, err = conn.Write(buf[:n])
		if err != nil {
			h.Close(conn)
			conn.Close()
			return
		}
	}
}

func NewHandler(ctx context.Context, instance *vcore.Instance) core.ConnectionHandler {
	h := &handler{
		ctx:        ctx,
		v:          instance,
		conns:      make(map[core.Connection]*connEntry, 16),
		dnsRespCh:  make(chan *dnsRespEntry, 1024),
		dispatched: make(map[core.Connection]bool, 16),
		dnsClient:  instance.GetFeature(vdns.ClientType()).(vdns.Client),
	}
	go h.handleDNSResponse()
	return h
}

func (h *handler) Connect(conn core.Connection, target net.Addr) error {
	dest := vnet.DestinationFromAddr(target)
	sid := vsession.NewID()
	ctx := vsession.ContextWithID(h.ctx, sid)
	c, err := vcore.Dial(ctx, h.v, dest)
	if err != nil {
		return errors.New(fmt.Sprintf("dial V proxy connection failed: %v", err))
	}
	// Note that ctx here is used for canceling fetching input goroutine, not
	// canceling the connection, thus create the cancelable context after Dial().
	ctx, cancel := context.WithCancel(ctx)
	h.Lock()
	h.conns[conn] = &connEntry{
		conn:                c,
		target:              dest,
		cancelFetchingInput: cancel,
		fetchingInputCtx:    ctx,
	}
	h.Unlock()
	go h.fetchInput(conn)
	log.Printf("new proxy connection for target: %s:%s", target.Network(), target.String())
	return nil
}

func (h *handler) DidReceive(conn core.Connection, data []byte) error {
	h.Lock()
	c, ok := h.conns[conn]
	done, ok2 := h.dispatched[conn]
	h.Unlock()
	if ok2 && done {
		// Request already dispatched to V2Ray, ignore.
		return nil
	}
	if ok {
		// If it's a DNS request of type A or AAAA and has only one question,
		// handle it with V2Ray's DNS client, otherwise treat as normal TCP/UDP
		// traffic.
		if c.target.Network == vnet.Network_UDP &&
			c.target.Port.Value() == proxy.COMMON_DNS_PORT &&
			h.shouldAcceptDNSQuery(data) {

			// Parse DNS request and hand to V2Ray, upon V2Ray returns []net.IP,
			// packing them into dns.Msg response message and write back to the client.
			go h.handleDNSQuery(conn, append([]byte(nil), data...))

			// The DNS request has passed to V2Ray for handling, we are safe
			// to cancel the fetching goroutine, but be careful do not close the
			// connection.
			c.cancelFetchingInput()

			h.Lock()
			// The request is successfully handed to V2Ray, mark as dispatched so
			// subsequent retransmissions should be ignored.
			h.dispatched[conn] = true
			h.Unlock()
		} else {
			_, err := c.conn.Write(data)
			if err != nil {
				h.Close(conn)
				return errors.New(fmt.Sprintf("write remote failed: %v", err))
			}
		}
		return nil
	} else {
		h.Close(conn)
		return errors.New(fmt.Sprintf("proxy connection %v->%v does not exists", conn.LocalAddr(), conn.RemoteAddr()))
	}
}

func (h *handler) DidSend(conn core.Connection, len uint16) {
	// unused
}

func (h *handler) DidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) LocalDidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) Close(conn core.Connection) {
	h.Lock()
	defer h.Unlock()

	if c, found := h.conns[conn]; found {
		c.conn.Close()
	}
	delete(h.conns, conn)
}
<commit_msg>fix a misbehavior against aaaa dns response<commit_after>package v2ray

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"
	"time"

	"github.com/miekg/dns"
	vcore "v2ray.com/core"
	vnet "v2ray.com/core/common/net"
	vsession "v2ray.com/core/common/session"
	vdns "v2ray.com/core/features/dns"

	"github.com/eycorsican/go-tun2socks/core"
	"github.com/eycorsican/go-tun2socks/proxy"
)

func isIPv4(ip net.IP) bool {
	if ip.To4() != nil {
		return true
	}
	return false
}

func isIPv6(ip net.IP) bool {
	// To16() also valid for ipv4, ensure it's not an ipv4 address
	if ip.To4() != nil {
		return false
	}
	if ip.To16() != nil {
		return true
	}
	return false
}

type connEntry struct {
	conn                net.Conn
	target              vnet.Destination
	cancelFetchingInput context.CancelFunc
	fetchingInputCtx    context.Context
}

type dnsRespEntry struct {
	err  error
	data []byte
	conn core.Connection
}

type handler struct {
	sync.Mutex

	ctx        context.Context
	v          *vcore.Instance
	conns      map[core.Connection]*connEntry
	dispatched map[core.Connection]bool
	dnsRespCh  chan *dnsRespEntry
	dnsClient  vdns.Client
}

func (h *handler) shouldAcceptDNSQuery(data []byte) bool {
	req := new(dns.Msg)
	err := req.Unpack(data)
	if err != nil {
		return false
	}

	// TODO: allow multiple question
	if len(req.Question) != 1 {
		return false
	}

	qtype := req.Question[0].Qtype
	if qtype != dns.TypeA && qtype != dns.TypeAAAA {
		return false
	}

	qclass := req.Question[0].Qclass
	if qclass != dns.ClassINET {
		return false
	}

	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	if _, ok := dns.IsDomainName(domain); !ok {
		return false
	}

	return true
}

func (h *handler) handleDNSQuery(conn core.Connection, data []byte) {
	var err error
	var answer []byte = nil
	defer func() {
		h.dnsRespCh <- &dnsRespEntry{conn: conn, data: answer, err: err}
	}()

	// No error checks here because they are already done in shouldAcceptDNSQuery()
	req := new(dns.Msg)
	req.Unpack(data)
	qtype := req.Question[0].Qtype
	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	log.Printf("dispatch dns request for domain: %v (%v)", domain, qtype)
	ips, err := h.dnsClient.LookupIP(domain)
	if err != nil {
		err = errors.New(fmt.Sprintf("lookup ip failed: %v", err))
		return
	}

	resp := new(dns.Msg)
	resp = resp.SetReply(req)
	resp.RecursionAvailable = true
	for _, ip := range ips {
		if isIPv4(ip) && qtype == dns.TypeA {
			resp.Answer = append(resp.Answer, &dns.A{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 4,
				},
				A: ip,
			})
		} else if isIPv6(ip) && qtype == dns.TypeAAAA {
			resp.Answer = append(resp.Answer, &dns.AAAA{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeAAAA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 16,
				},
				AAAA: ip,
			})
		}
	}
	if len(resp.Answer) == 0 {
		// Has A records but no wanted AAAA records
		if qtype == dns.TypeAAAA && len(ips) != 0 {
			// https://tools.ietf.org/html/rfc4074#section-3
			resp = resp.SetRcode(req, dns.RcodeSuccess)
		} else {
			err = errors.New(fmt.Sprintf("no answer for %v (%v) (%v)", domain, qtype, len(ips)))
			return
		}
	}
	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)
	dnsAnswer, err := resp.PackBuffer(buf)
	if err != nil {
		err = errors.New(fmt.Sprintf("packing dns resp msg failed: %v", err))
	} else {
		answer = append([]byte(nil), dnsAnswer...)
	}
	return
}

func (h *handler) handleDNSResponse() {
	for {
		select {
		case respEntry := <-h.dnsRespCh:
			if respEntry.err == nil {
				_, err := respEntry.conn.Write(respEntry.data)
				if err != nil {
					log.Printf("write dns response to local failed: %v", err)
				}
			} else {
				log.Printf("dispatch dns request failed: %v", respEntry.err)
			}
			h.Close(respEntry.conn)
			respEntry.conn.Close()
		}
	}
}

func (h *handler) fetchInput(conn core.Connection) {
	h.Lock()
	c, ok := h.conns[conn]
	h.Unlock()
	if !ok {
		return
	}

	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)

FetchingLoop:
	for {
		c.conn.SetReadDeadline(time.Now().Add(4 * time.Second))
		n, err := c.conn.Read(buf)
		if err, ok := err.(net.Error); ok && err.Timeout() {
			select {
			case <-c.fetchingInputCtx.Done():
				// Request was handed to V2Ray, stop fetching but leave the
				// connection open.
				return
			default:
				continue FetchingLoop
			}
		}
		if err != nil {
			h.Close(conn)
			conn.Close()
			return
		}
		_, err = conn.Write(buf[:n])
		if err != nil {
			h.Close(conn)
			conn.Close()
			return
		}
	}
}

func NewHandler(ctx context.Context, instance *vcore.Instance) core.ConnectionHandler {
	h := &handler{
		ctx:        ctx,
		v:          instance,
		conns:      make(map[core.Connection]*connEntry, 16),
		dnsRespCh:  make(chan *dnsRespEntry, 1024),
		dispatched: make(map[core.Connection]bool, 16),
		dnsClient:  instance.GetFeature(vdns.ClientType()).(vdns.Client),
	}
	go h.handleDNSResponse()
	return h
}

func (h *handler) Connect(conn core.Connection, target net.Addr) error {
	dest := vnet.DestinationFromAddr(target)
	sid := vsession.NewID()
	ctx := vsession.ContextWithID(h.ctx, sid)
	c, err := vcore.Dial(ctx, h.v, dest)
	if err != nil {
		return errors.New(fmt.Sprintf("dial V proxy connection failed: %v", err))
	}
	// Note that ctx here is used for canceling fetching input goroutine, not
	// canceling the connection, thus create the cancelable context after Dial().
	ctx, cancel := context.WithCancel(ctx)
	h.Lock()
	h.conns[conn] = &connEntry{
		conn:                c,
		target:              dest,
		cancelFetchingInput: cancel,
		fetchingInputCtx:    ctx,
	}
	h.Unlock()
	go h.fetchInput(conn)
	log.Printf("new proxy connection for target: %s:%s", target.Network(), target.String())
	return nil
}

func (h *handler) DidReceive(conn core.Connection, data []byte) error {
	h.Lock()
	c, ok := h.conns[conn]
	done, ok2 := h.dispatched[conn]
	h.Unlock()
	if ok2 && done {
		// Request already dispatched to V2Ray, ignore.
		return nil
	}
	if ok {
		// If it's a DNS request of type A or AAAA and has only one question,
		// handle it with V2Ray's DNS client, otherwise treat as normal TCP/UDP
		// traffic.
		if c.target.Network == vnet.Network_UDP &&
			c.target.Port.Value() == proxy.COMMON_DNS_PORT &&
			h.shouldAcceptDNSQuery(data) {

			// Parse DNS request and hand to V2Ray, upon V2Ray returns []net.IP,
			// packing them into dns.Msg response message and write back to the client.
			go h.handleDNSQuery(conn, append([]byte(nil), data...))

			// The DNS request has passed to V2Ray for handling, we are safe
			// to cancel the fetching goroutine, but be careful do not close the
			// connection.
			c.cancelFetchingInput()

			h.Lock()
			// The request is successfully handed to V2Ray, mark as dispatched so
			// subsequent retransmissions should be ignored.
			h.dispatched[conn] = true
			h.Unlock()
		} else {
			_, err := c.conn.Write(data)
			if err != nil {
				h.Close(conn)
				return errors.New(fmt.Sprintf("write remote failed: %v", err))
			}
		}
		return nil
	} else {
		h.Close(conn)
		return errors.New(fmt.Sprintf("proxy connection %v->%v does not exists", conn.LocalAddr(), conn.RemoteAddr()))
	}
}

func (h *handler) DidSend(conn core.Connection, len uint16) {
	// unused
}

func (h *handler) DidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) LocalDidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) Close(conn core.Connection) {
	h.Lock()
	defer h.Unlock()

	if c, found := h.conns[conn]; found {
		c.conn.Close()
	}
	delete(h.conns, conn)
}
<|endoftext|>


<commit_before>package v2ray

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"

	"github.com/miekg/dns"
	vcore "v2ray.com/core"
	vnet "v2ray.com/core/common/net"
	vsession "v2ray.com/core/common/session"
	vdns "v2ray.com/core/features/dns"

	"github.com/eycorsican/go-tun2socks/core"
	"github.com/eycorsican/go-tun2socks/proxy"
)

func isIPv4(ip net.IP) bool {
	if ip.To4() != nil {
		return true
	}
	return false
}

func isIPv6(ip net.IP) bool {
	// To16() also valid for ipv4, ensure it's not an ipv4 address
	if ip.To4() != nil {
		return false
	}
	if ip.To16() != nil {
		return true
	}
	return false
}

type connEntry struct {
	conn                net.Conn
	target              vnet.Destination
	cancelFetchingInput context.CancelFunc
	fetchingInputCtx    context.Context
}

type dnsRespEntry struct {
	err  error
	data []byte
	conn core.Connection
}

type handler struct {
	sync.Mutex

	ctx        context.Context
	v          *vcore.Instance
	conns      map[core.Connection]*connEntry
	dispatched map[core.Connection]bool
	dnsRespCh  chan *dnsRespEntry
	dnsClient  vdns.Client
}

func (h *handler) shouldAcceptDNSQuery(data []byte) bool {
	req := new(dns.Msg)
	err := req.Unpack(data)
	if err != nil {
		return false
	}

	// TODO: allow multiple question
	if len(req.Question) != 1 {
		return false
	}

	qtype := req.Question[0].Qtype
	if qtype != dns.TypeA && qtype != dns.TypeAAAA {
		return false
	}

	qclass := req.Question[0].Qclass
	if qclass != dns.ClassINET {
		return false
	}

	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	if _, ok := dns.IsDomainName(domain); !ok {
		return false
	}

	return true
}

func (h *handler) handleDNSQuery(conn core.Connection, data []byte) {
	var err error
	var answer []byte = nil
	defer func() {
		h.dnsRespCh <- &dnsRespEntry{conn: conn, data: answer, err: err}
	}()

	// No error checks here because they are already done in shouldAcceptDNSQuery()
	req := new(dns.Msg)
	req.Unpack(data)
	qtype := req.Question[0].Qtype
	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	log.Printf("dispatch dns request for domain: %v (%v)", domain, qtype)
	ips, err := h.dnsClient.LookupIP(domain)
	if err != nil {
		err = errors.New(fmt.Sprintf("lookup ip failed: %v", err))
		return
	}

	resp := new(dns.Msg)
	resp = resp.SetReply(req)
	resp.RecursionAvailable = true
	for _, ip := range ips {
		if isIPv4(ip) && qtype == dns.TypeA {
			resp.Answer = append(resp.Answer, &dns.A{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 4,
				},
				A: ip,
			})
		} else if isIPv6(ip) && qtype == dns.TypeAAAA {
			resp.Answer = append(resp.Answer, &dns.AAAA{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeAAAA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 16,
				},
				AAAA: ip,
			})
		}
	}
	if len(resp.Answer) == 0 {
		// Has A records but no wanted AAAA records
		if qtype == dns.TypeAAAA && len(ips) != 0 {
			// https://tools.ietf.org/html/rfc4074#section-3
			resp = resp.SetRcode(req, dns.RcodeSuccess)
		} else {
			err = errors.New(fmt.Sprintf("no answer for %v (%v) (%v)", domain, qtype, len(ips)))
			return
		}
	}
	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)
	dnsAnswer, err := resp.PackBuffer(buf)
	if err != nil {
		err = errors.New(fmt.Sprintf("packing dns resp msg failed: %v", err))
	} else {
		answer = append([]byte(nil), dnsAnswer...)
	}
	return
}

func (h *handler) handleDNSResponse() {
	for {
		select {
		case respEntry := <-h.dnsRespCh:
			if respEntry.err == nil {
				_, err := respEntry.conn.Write(respEntry.data)
				if err != nil {
					log.Printf("write dns response to local failed: %v", err)
				}
			} else {
				log.Printf("dispatch dns request failed: %v", respEntry.err)
			}
			h.Close(respEntry.conn)
			respEntry.conn.Close()
		}
	}
}

func (h *handler) fetchInput(conn core.Connection) {
	h.Lock()
	c, ok := h.conns[conn]
	h.Unlock()
	if !ok {
		return
	}

	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)

	for {
		n, err := c.conn.Read(buf)
		select {
		case <-c.fetchingInputCtx.Done():
			// Request was re-dipatched to V2Ray, stop fetching but leave the
			// connection open.
			return
		default:
		}
		select {
		case <-c.fetchingInputCtx.Done():
			// Request was re-dipatched to V2Ray, stop fetching but leave the
			// connection open.
			return
		default:
			if err != nil && n <= 0 {
				h.Close(conn)
				conn.Close()
				return
			}
			_, err = conn.Write(buf[:n])
			if err != nil {
				h.Close(conn)
				conn.Close()
				return
			}
		}
	}
}

func NewHandler(ctx context.Context, instance *vcore.Instance) core.ConnectionHandler {
	h := &handler{
		ctx:        ctx,
		v:          instance,
		conns:      make(map[core.Connection]*connEntry, 16),
		dnsRespCh:  make(chan *dnsRespEntry, 1024),
		dispatched: make(map[core.Connection]bool, 16),
		dnsClient:  instance.GetFeature(vdns.ClientType()).(vdns.Client),
	}
	go h.handleDNSResponse()
	return h
}

func (h *handler) Connect(conn core.Connection, target net.Addr) error {
	dest := vnet.DestinationFromAddr(target)
	sid := vsession.NewID()
	ctx := vsession.ContextWithID(h.ctx, sid)
	c, err := vcore.Dial(ctx, h.v, dest)
	if err != nil {
		return errors.New(fmt.Sprintf("dial V proxy connection failed: %v", err))
	}
	// Note that ctx here is used for canceling fetching input goroutine, not
	// canceling the connection, thus create the cancelable context after Dial().
	ctx, cancel := context.WithCancel(ctx)
	h.Lock()
	h.conns[conn] = &connEntry{
		conn:                c,
		target:              dest,
		cancelFetchingInput: cancel,
		fetchingInputCtx:    ctx,
	}
	h.Unlock()
	go h.fetchInput(conn)
	log.Printf("new proxy connection for target: %s:%s", target.Network(), target.String())
	return nil
}

func (h *handler) DidReceive(conn core.Connection, data []byte) error {
	h.Lock()
	c, ok := h.conns[conn]
	done, ok2 := h.dispatched[conn]
	h.Unlock()
	if ok2 && done {
		// Request already dispatched to V2Ray, ignore.
		return nil
	}
	if ok {
		// If it's a DNS request of type A or AAAA and has only one question,
		// handle it with V2Ray's DNS client, otherwise treat as normal TCP/UDP
		// traffic.
		if c.target.Network == vnet.Network_UDP &&
			c.target.Port.Value() == proxy.COMMON_DNS_PORT &&
			h.shouldAcceptDNSQuery(data) {

			// Parse DNS request and hand to V2Ray, upon V2Ray returns []net.IP,
			// packing them into dns.Msg response message and write back to the client.
			go h.handleDNSQuery(conn, append([]byte(nil), data...))

			// The DNS request has passed to V2Ray for handling, we are safe
			// to cancel the fetching goroutine, but be careful do not close the
			// connection.
			c.cancelFetchingInput()
			c.conn.Close()

			h.Lock()
			// The request is successfully handed to V2Ray, mark as dispatched so
			// subsequent retransmissions should be ignored.
			h.dispatched[conn] = true
			h.Unlock()
		} else {
			_, err := c.conn.Write(data)
			if err != nil {
				h.Close(conn)
				return errors.New(fmt.Sprintf("write remote failed: %v", err))
			}
		}
		return nil
	} else {
		h.Close(conn)
		return errors.New(fmt.Sprintf("proxy connection %v->%v does not exists", conn.LocalAddr(), conn.RemoteAddr()))
	}
}

func (h *handler) DidSend(conn core.Connection, len uint16) {
	// unused
}

func (h *handler) DidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) LocalDidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) Close(conn core.Connection) {
	h.Lock()
	defer h.Unlock()

	if c, found := h.conns[conn]; found {
		c.conn.Close()
	}
	delete(h.conns, conn)
}
<commit_msg>specify dns query type<commit_after>package v2ray

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net"
	"sync"

	"github.com/miekg/dns"
	vcore "v2ray.com/core"
	vnet "v2ray.com/core/common/net"
	vsession "v2ray.com/core/common/session"
	vdns "v2ray.com/core/features/dns"

	"github.com/eycorsican/go-tun2socks/core"
	"github.com/eycorsican/go-tun2socks/proxy"
)

func isIPv4(ip net.IP) bool {
	if ip.To4() != nil {
		return true
	}
	return false
}

func isIPv6(ip net.IP) bool {
	// To16() also valid for ipv4, ensure it's not an ipv4 address
	if isIPv4(ip) {
		return false
	}
	if ip.To16() != nil {
		return true
	}
	return false
}

type connEntry struct {
	conn                net.Conn
	target              vnet.Destination
	cancelFetchingInput context.CancelFunc
	fetchingInputCtx    context.Context
}

type dnsRespEntry struct {
	err  error
	data []byte
	conn core.Connection
}

type handler struct {
	sync.Mutex

	ctx        context.Context
	v          *vcore.Instance
	conns      map[core.Connection]*connEntry
	dispatched map[core.Connection]bool
	dnsRespCh  chan *dnsRespEntry
	dnsClient  vdns.Client
}

func (h *handler) shouldAcceptDNSQuery(data []byte) bool {
	req := new(dns.Msg)
	err := req.Unpack(data)
	if err != nil {
		return false
	}

	// TODO: allow multiple question
	if len(req.Question) != 1 {
		return false
	}

	qtype := req.Question[0].Qtype
	if qtype != dns.TypeA && qtype != dns.TypeAAAA {
		return false
	}

	qclass := req.Question[0].Qclass
	if qclass != dns.ClassINET {
		return false
	}

	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	if _, ok := dns.IsDomainName(domain); !ok {
		return false
	}

	return true
}

func (h *handler) handleDNSQuery(conn core.Connection, data []byte) {
	var err error
	var answer []byte = nil
	defer func() {
		h.dnsRespCh <- &dnsRespEntry{conn: conn, data: answer, err: err}
	}()

	// No error checks here because they are already done in shouldAcceptDNSQuery()
	req := new(dns.Msg)
	req.Unpack(data)
	qtype := req.Question[0].Qtype
	fqdn := req.Question[0].Name
	domain := fqdn[:len(fqdn)-1]

	log.Printf("dispatch dns request for domain: %v (%v)", domain, qtype)
	var ips []net.IP
	switch qtype {
	case dns.TypeA:
		if dnsClient, ok := h.dnsClient.(vdns.IPv4Lookup); ok {
			ips, err = dnsClient.LookupIPv4(domain)
		} else {
			ips, err = h.dnsClient.LookupIP(domain)
		}
	case dns.TypeAAAA:
		ips, err = h.dnsClient.LookupIP(domain)
	default:
		err = errors.New(fmt.Sprintf("impossible qtype: %v", qtype))
		return
	}
	if err != nil {
		err = errors.New(fmt.Sprintf("lookup ip failed: %v", err))
		return
	}

	resp := new(dns.Msg)
	resp = resp.SetReply(req)
	resp.RecursionAvailable = true
	for _, ip := range ips {
		if qtype == dns.TypeA && isIPv4(ip) {
			resp.Answer = append(resp.Answer, &dns.A{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 4,
				},
				A: ip,
			})
		} else if qtype == dns.TypeAAAA && isIPv6(ip) {
			resp.Answer = append(resp.Answer, &dns.AAAA{
				Hdr: dns.RR_Header{
					Name:     fqdn,
					Rrtype:   dns.TypeAAAA,
					Class:    dns.ClassINET,
					Ttl:      150, // cached in V2Ray
					Rdlength: 16,
				},
				AAAA: ip,
			})
		}
	}
	if len(resp.Answer) == 0 {
		// Has A records but no wanted AAAA records
		if qtype == dns.TypeAAAA && len(ips) != 0 {
			// https://tools.ietf.org/html/rfc4074#section-3
			resp = resp.SetRcode(req, dns.RcodeSuccess)
		} else {
			err = errors.New(fmt.Sprintf("no answer for %v (%v) (%v)", domain, qtype, len(ips)))
			return
		}
	}
	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)
	dnsAnswer, err := resp.PackBuffer(buf)
	if err != nil {
		err = errors.New(fmt.Sprintf("packing dns resp msg failed: %v", err))
	} else {
		answer = append([]byte(nil), dnsAnswer...)
	}
	return
}

func (h *handler) handleDNSResponse() {
	for {
		select {
		case respEntry := <-h.dnsRespCh:
			if respEntry.err == nil {
				_, err := respEntry.conn.Write(respEntry.data)
				if err != nil {
					log.Printf("write dns response to local failed: %v", err)
				}
			} else {
				log.Printf("dispatch dns request failed: %v", respEntry.err)
			}
			h.Close(respEntry.conn)
			respEntry.conn.Close()
		}
	}
}

func (h *handler) fetchInput(conn core.Connection) {
	h.Lock()
	c, ok := h.conns[conn]
	h.Unlock()
	if !ok {
		return
	}

	buf := core.NewBytes(core.BufSize)
	defer core.FreeBytes(buf)

	for {
		n, err := c.conn.Read(buf)
		select {
		case <-c.fetchingInputCtx.Done():
			// Request was re-dipatched to V2Ray, stop fetching but leave the
			// connection open.
			return
		default:
		}
		select {
		case <-c.fetchingInputCtx.Done():
			// Request was re-dipatched to V2Ray, stop fetching but leave the
			// connection open.
			return
		default:
			if err != nil && n <= 0 {
				h.Close(conn)
				conn.Close()
				return
			}
			_, err = conn.Write(buf[:n])
			if err != nil {
				h.Close(conn)
				conn.Close()
				return
			}
		}
	}
}

func NewHandler(ctx context.Context, instance *vcore.Instance) core.ConnectionHandler {
	h := &handler{
		ctx:        ctx,
		v:          instance,
		conns:      make(map[core.Connection]*connEntry, 16),
		dnsRespCh:  make(chan *dnsRespEntry, 1024),
		dispatched: make(map[core.Connection]bool, 16),
		dnsClient:  instance.GetFeature(vdns.ClientType()).(vdns.Client),
	}
	go h.handleDNSResponse()
	return h
}

func (h *handler) Connect(conn core.Connection, target net.Addr) error {
	dest := vnet.DestinationFromAddr(target)
	sid := vsession.NewID()
	ctx := vsession.ContextWithID(h.ctx, sid)
	c, err := vcore.Dial(ctx, h.v, dest)
	if err != nil {
		return errors.New(fmt.Sprintf("dial V proxy connection failed: %v", err))
	}
	// Note that ctx here is used for canceling fetching input goroutine, not
	// canceling the connection, thus create the cancelable context after Dial().
	ctx, cancel := context.WithCancel(ctx)
	h.Lock()
	h.conns[conn] = &connEntry{
		conn:                c,
		target:              dest,
		cancelFetchingInput: cancel,
		fetchingInputCtx:    ctx,
	}
	h.Unlock()
	go h.fetchInput(conn)
	log.Printf("new proxy connection for target: %s:%s", target.Network(), target.String())
	return nil
}

func (h *handler) DidReceive(conn core.Connection, data []byte) error {
	h.Lock()
	c, ok := h.conns[conn]
	done, ok2 := h.dispatched[conn]
	h.Unlock()
	if ok2 && done {
		// Request already dispatched to V2Ray, ignore.
		return nil
	}
	if ok {
		// If it's a DNS request of type A or AAAA and has only one question,
		// handle it with V2Ray's DNS client, otherwise treat as normal TCP/UDP
		// traffic.
		if c.target.Network == vnet.Network_UDP &&
			c.target.Port.Value() == proxy.COMMON_DNS_PORT &&
			h.shouldAcceptDNSQuery(data) {

			// Parse DNS request and hand to V2Ray, upon V2Ray returns []net.IP,
			// packing them into dns.Msg response message and write back to the client.
			go h.handleDNSQuery(conn, append([]byte(nil), data...))

			// The DNS request has passed to V2Ray for handling, we are safe
			// to cancel the fetching goroutine, but be careful do not close the
			// connection.
			c.cancelFetchingInput()
			c.conn.Close()

			h.Lock()
			// The request is successfully handed to V2Ray, mark as dispatched so
			// subsequent retransmissions should be ignored.
			h.dispatched[conn] = true
			h.Unlock()
		} else {
			_, err := c.conn.Write(data)
			if err != nil {
				h.Close(conn)
				return errors.New(fmt.Sprintf("write remote failed: %v", err))
			}
		}
		return nil
	} else {
		h.Close(conn)
		return errors.New(fmt.Sprintf("proxy connection %v->%v does not exists", conn.LocalAddr(), conn.RemoteAddr()))
	}
}

func (h *handler) DidSend(conn core.Connection, len uint16) {
	// unused
}

func (h *handler) DidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) LocalDidClose(conn core.Connection) {
	h.Close(conn)
}

func (h *handler) Close(conn core.Connection) {
	h.Lock()
	defer h.Unlock()

	if c, found := h.conns[conn]; found {
		c.conn.Close()
	}
	delete(h.conns, conn)
}
<|endoftext|>