// RegisterTmp create a ephemeral node, and watch it, if node droped then send a SIGQUIT to self.
func RegisterTemp(conn *zk.Conn, fpath string, data []byte) error {
tpath, err := conn.Create(path.Join(fpath)+"/", data, zk.FlagEphemeral|zk.FlagSequence, zk.WorldACL(zk.PermAll))
if err != nil {
log.Error("conn.Create(\"%s\", \"%s\", zk.FlagEphemeral|zk.FlagSequence) error(%v)", fpath, string(data), err)
return err
}
log.Debug("create a zookeeper node:%s", tpath)
// watch self
go func() {
for {
log.Info("zk path: \"%s\" set a watch", tpath)
exist, _, watch, err := conn.ExistsW(tpath)
if err != nil {
log.Error("zk.ExistsW(\"%s\") error(%v)", tpath, err)
log.Warn("zk path: \"%s\" set watch failed, kill itself", tpath)
killSelf()
return
}
if !exist {
log.Warn("zk path: \"%s\" not exist, kill itself", tpath)
killSelf()
return
}
event := <-watch
log.Info("zk path: \"%s\" receive a event %v", tpath, event)
}
}()
return nil
}
func main() {
flag.BoolVar(&DEBUG, "d", false, "debug mode")
mode := flag.String("m", "local", "mode: local or remote")
laddr := flag.String("l", "127.0.0.1:9000", "local address")
raddr := flag.String("r", "127.0.0.1:9001", "remote address")
cert := flag.String("c", "", "certificate (remote)")
key := flag.String("k", "", "private key (remote)")
tcount := flag.Int("t", 1, "tunnel count (local)")
flag.Parse()
// go func() {
// var addr string
// if *mode == "local" {
// addr = "localhost:6060"
// } else {
// addr = "localhost:6061"
// }
// err := http.ListenAndServe(addr, nil)
// if err != nil {
// log.Error("pprof error:%v", err)
// }
// }()
if DEBUG {
log.AddFilter("stdout", log.DEBUG, log.NewConsoleLogWriter())
} else {
log.AddFilter("stdout", log.INFO, log.NewConsoleLogWriter())
}
switch *mode {
case "local":
localServer, err := NewLocalServer(*laddr, *raddr, *tcount)
if err != nil {
log.Error("NewLocalServer error:%v", err)
break
}
localServer.Serve()
case "remote":
remoteServer, err := NewRemoteServer(*laddr, *raddr)
if err != nil {
log.Error("NewRemoteServer error:%v", err)
break
}
remoteServer.Serve(*cert, *key)
}
log.Info("exit")
time.Sleep(2 * time.Second)
}
func (rs *RemoteServer) serveStream(stream net.Conn) {
defer stream.Close()
start := time.Now()
conn, err := net.Dial("tcp", rs.raddr)
if err != nil {
log.Error("connect to remote error:%v", err)
return
}
defer conn.Close()
atomic.AddInt32(&rs.streamCount, 1)
streamID := atomic.AddUint64(&rs.streamID, 1)
readChan := make(chan int64, 1)
writeChan := make(chan int64, 1)
var readBytes int64
var writeBytes int64
go pipe(stream, conn, readChan)
go pipe(conn, stream, writeChan)
for i := 0; i < 2; i++ {
select {
case readBytes = <-readChan:
stream.Close()
log.Debug("[#%d] read %d bytes", streamID, readBytes)
case writeBytes = <-writeChan:
// DON'T call conn.Close, it will trigger an error in pipe.
// Just close stream, let the conn copy finished normally
log.Debug("[#%d] write %d bytes", streamID, writeBytes)
}
}
log.Info("[#%d] r:%d w:%d t:%v c:%d",
streamID, readBytes, writeBytes, time.Now().Sub(start), atomic.LoadInt32(&rs.streamCount))
atomic.AddInt32(&rs.streamCount, -1)
}
func (ls *LocalServer) transport(conn net.Conn) {
defer conn.Close()
start := time.Now()
stream, err := ls.openStream()
if err != nil {
log.Error("open stream for %s error:%v", conn.RemoteAddr().String(), err)
return
}
defer stream.Close()
atomic.AddInt32(&ls.streamCount, 1)
streamID := atomic.AddUint64(&ls.streamID, 1)
readChan := make(chan int64, 1)
writeChan := make(chan int64, 1)
var readBytes int64
var writeBytes int64
go pipe(conn, stream, readChan)
go pipe(stream, conn, writeChan)
for i := 0; i < 2; i++ {
select {
case readBytes = <-readChan:
// DON'T call conn.Close, it will trigger an error in pipe.
// Just close stream, let the conn copy finished normally
log.Debug("[#%d] read %d bytes", streamID, readBytes)
case writeBytes = <-writeChan:
stream.Close()
log.Debug("[#%d] write %d bytes", streamID, writeBytes)
}
}
log.Info("[#%d] r:%d w:%d t:%v c:%d",
streamID, readBytes, writeBytes, time.Now().Sub(start), atomic.LoadInt32(&ls.streamCount))
atomic.AddInt32(&ls.streamCount, -1)
}
func (ls *LocalServer) Serve() {
l, err := net.Listen("tcp", ls.laddr)
if err != nil {
log.Error("listen [%s] error:%v", ls.laddr, err)
return
}
for {
c, err := l.Accept()
if err != nil {
log.Error("accept connection error:%v", err)
continue
}
go ls.transport(c)
}
}
// Connect connect to zookeeper, and start a goroutine log the event.
func Connect(addr []string, timeout time.Duration) (*zk.Conn, error) {
conn, session, err := zk.Connect(addr, timeout)
if err != nil {
log.Error("zk.Connect(\"%v\", %d) error(%v)", addr, timeout, err)
return nil, err
}
go func() {
for {
event := <-session
log.Debug("zookeeper get a event: %s", event.State.String())
}
}()
return conn, nil
}
func (rs *RemoteServer) Serve(certFile, keyFile string) {
cert, err := tls.LoadX509KeyPair(certFile, keyFile)
if err != nil {
log.Error("load cert/key error:%v", err)
return
}
config := tls.Config{
Certificates: []tls.Certificate{cert},
}
l, err := tls.Listen("tcp", rs.laddr, &config)
if err != nil {
log.Error("net.ListenTCP(%s) error:%v", rs.laddr, err)
return
}
for {
stream, err := l.Accept()
if err != nil {
log.Error("listenner accept connection error:%v", err)
continue
}
go rs.serveStream(stream)
}
}
// GetNodes get all child from zk path.
func GetNodes(conn *zk.Conn, path string) ([]string, error) {
nodes, stat, err := conn.Children(path)
if err != nil {
if err == zk.ErrNoNode {
return nil, ErrNodeNotExist
}
log.Error("zk.Children(\"%s\") error(%v)", path, err)
return nil, err
}
if stat == nil {
return nil, ErrNodeNotExist
}
if len(nodes) == 0 {
return nil, ErrNoChild
}
return nodes, nil
}
// Create create zookeeper path, if path exists ignore error
func Create(conn *zk.Conn, fpath string) error {
// create zk root path
tpath := ""
for _, str := range strings.Split(fpath, "/")[1:] {
tpath = path.Join(tpath, "/", str)
log.Debug("create zookeeper path: \"%s\"", tpath)
_, err := conn.Create(tpath, []byte(""), 0, zk.WorldACL(zk.PermAll))
if err != nil {
if err == zk.ErrNodeExists {
log.Warn("zk.create(\"%s\") exists", tpath)
} else {
log.Error("zk.create(\"%s\") error(%v)", tpath, err)
return err
}
}
}
return nil
}
// killSelf send a SIGQUIT to self.
func killSelf() {
if err := syscall.Kill(os.Getpid(), syscall.SIGQUIT); err != nil {
log.Error("syscall.Kill(%d, SIGQUIT) error(%v)", os.Getpid(), err)
}
}
