Permalink
Browse files

Merge "initialize bandwidth proxy"

  • Loading branch information...
2 parents 44fa668 + ca4155c commit 280a45744bc0797e88973f65847dc5d1a3146a64 @andyzh andyzh committed with Gerrit Code Review Oct 8, 2012
@@ -0,0 +1,74 @@
+package main
+
+import (
+ "flag"
+ "fmt"
+ "net"
+ "os"
+ "os/signal"
+ "proxy_tunnel/logger"
+ "proxy_tunnel/tunnel"
+ "time"
+)
+
+var logFile string // Log file name
+var ePort uint // External port proxy listen to
+var iIp string // Inner ip proxy connect to
+var iPort uint // Inner port proxy connect to
+var window uint // Time window to check the limit(in seconds)
+var limit uint64 // Limit size(in bytes) per time window include bosh inbound and outbound
+var signalChan chan os.Signal // Cahnel for signal
+
+func signalHand() {
+ signalChan = make(chan os.Signal, 1)
+ signal.Notify(signalChan, os.Interrupt, os.Kill)
+ select {
+ case <-signalChan:
+ tunnel.Stop()
+ os.Exit(-1)
+ }
+ return
+}
+
+func main() {
+ flag.UintVar(&ePort, "eport", 0, "port proxy listen")
+ flag.StringVar(&iIp, "iip", "127.0.0.1", "inner ip proxy connect to")
+ flag.UintVar(&iPort, "iport", 0, "inner port proxy connect to")
+ flag.Uint64Var(&limit, "limit", 0, "limit size per time window(in bytes)")
+ flag.UintVar(&window, "window", 0, "time window to check the limit(in seconds)")
+ flag.StringVar(&logFile, "l", "", "log file name")
+ flag.Usage = func() {
+ fmt.Fprintf(os.Stderr, "usage: %s -eport external_port -iport internal_port [-iip internal_ip] -limit limit_size -window time_window -l log_file\n", os.Args[0])
+ flag.PrintDefaults()
+ os.Exit(-1)
+ }
+ flag.Parse()
+ if ePort == 0 || iPort == 0 || logFile == "" || limit == 0 || window == 0 {
+ flag.Usage()
+ }
+ ip := net.ParseIP(iIp)
+ if ip == nil {
+ fmt.Fprintln(os.Stderr, "Invalid ip:", iIp)
+ flag.Usage()
+ }
+
+ err := logger.Init(logFile)
+ if err != nil {
+ fmt.Println("Init log file error:", err)
+ os.Exit(-1)
+ }
+ defer logger.Finalize()
+
+ // Handle signal
+ go signalHand()
+
+ t := tunnel.Tunnel{
+ EPort: ePort,
+ IIp: ip,
+ IPort: iPort,
+ Limit: limit,
+ Window: window,
+ CheckTime: time.Now(),
+ }
+ tunnel.Run(&t)
+}
@@ -0,0 +1,81 @@
+package logger
+
+import (
+ "fmt"
+ "log"
+ "os"
+ "runtime"
+)
+
+const (
+ ERR = 1
+ WARN = 2
+ DEBUG = 3
+ INFO = 4
+)
+
+type Logger struct {
+ ch chan string
+ file *os.File
+ fileName string
+}
+
+var obj *Logger
+
+func Init(logFile string) error {
+ if obj != nil {
+ return nil
+ }
+ obj = new(Logger)
+ if logFile != "" {
+ outFile, err := os.OpenFile(logFile, os.O_RDWR|os.O_APPEND|os.O_CREATE, 0666)
+ if err != nil {
+ return err
+ }
+ obj.file = outFile
+ }
+ if obj.file == nil {
+ obj.file = os.Stderr
+ }
+ obj.fileName = logFile
+ obj.ch = make(chan string, 100)
+ log.SetOutput(obj.file)
+ go Start()
+ return nil
+}
+
+func Start() {
+ for {
+ msg := <-obj.ch
+ if len(msg) == 0 {
+ runtime.Goexit()
+ }
+ log.Println(msg)
+ }
+}
+
+func Instance() *Logger {
+ return obj
+}
+
+var LogLevelMap = map[int8]string{
+ ERR: "ERROR",
+ WARN: "WARN",
+ DEBUG: "DEBUG",
+ INFO: "INFO",
+}
+
+func Log(level int8, fmtStr string, a ...interface{}) (err error) {
+ msg := fmt.Sprintf("[%s] %s", LogLevelMap[level], fmt.Sprintf(fmtStr, a...))
+ obj.ch <- msg
+ return nil
+}
+
+func Finalize() {
+ if obj != nil {
+ obj.ch <- ""
+ if obj.fileName != "" {
+ obj.file.Close()
+ }
+ }
+}
@@ -0,0 +1,70 @@
+package tunnel
+
+import (
+ "fmt"
+ "os"
+ "proxy_tunnel/logger"
+ "syscall"
+)
+
+var runTunnel *Tunnel
+var fdTunnelConn = make(map[int]*TunnelConn)
+var epollFd int
+
+func Stop() {
+ logger.Log(logger.INFO, "Stop proxy server")
+ if runTunnel != nil {
+ syscall.Close(runTunnel.LFd)
+ for _, tc := range fdTunnelConn {
+ syscall.Close(tc.EFd)
+ syscall.Close(tc.IFd)
+ }
+ }
+}
+
+func Run(t *Tunnel) {
+ var err error
+ epollFd, err = syscall.EpollCreate(1024)
+ if err != nil {
+ logger.Log(logger.ERR, "Create epoll fd error [%s]", err)
+ os.Exit(-2)
+ }
+
+ for _, step := range initStep {
+ err = step.Action(t)
+ if err != nil {
+ fmt.Fprintf(os.Stderr, step.ErrFmt, err)
+ os.Exit(-2)
+ }
+ }
+ runTunnel = t
+
+ events := make([]syscall.EpollEvent, 10, 10)
+ for {
+ en, err := syscall.EpollWait(epollFd, events, 1000)
+ if err != nil {
+ logger.Log(logger.ERR, "Wail epoll fd error [%s]", err)
+ os.Exit(-2)
+ }
+ for i := 0; i < en; i++ {
+ ee := events[i]
+ if runTunnel.LFd == int(ee.Fd) {
+ runTunnel.newConn()
+ continue
+ }
+ tc, ok := fdTunnelConn[int(ee.Fd)]
+ if !ok {
+ continue
+ }
+ if ee.Events&syscall.EPOLLIN != 0 {
+ tc.handleIn(int(ee.Fd))
+ }
+ if ee.Events&syscall.EPOLLOUT != 0 {
+ tc.handleOut(int(ee.Fd))
+ }
+ if ee.Events&syscall.EPOLLHUP != 0 {
+ tc.shutdown()
+ }
+ }
+ }
+}
Oops, something went wrong.

0 comments on commit 280a457

Please sign in to comment.