Support Windows service with go-svc
parent
bba0bb6613
commit
f6449feb55
|
@ -0,0 +1,6 @@
|
||||||
|
module "github.com/productionwentdown/forward"
|
||||||
|
|
||||||
|
require (
|
||||||
|
"github.com/judwhite/go-svc" v1.0.0
|
||||||
|
"golang.org/x/sys" v0.0.0-20180322165403-91ee8cde4354
|
||||||
|
)
|
100
main.go
100
main.go
|
@ -1,68 +1,66 @@
|
||||||
package main
|
package main
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"flag"
|
|
||||||
"io"
|
|
||||||
"log"
|
"log"
|
||||||
"net"
|
"os"
|
||||||
|
"path/filepath"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"github.com/judwhite/go-svc/svc"
|
||||||
)
|
)
|
||||||
|
|
||||||
var listen string
|
type program struct {
|
||||||
var connect string
|
logFile *os.File
|
||||||
|
wg sync.WaitGroup
|
||||||
|
}
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
flag.StringVar(&listen, "listen", ":8000", "listen on ip and port")
|
prg := &program{}
|
||||||
flag.StringVar(&connect, "connect", "", "forward to ip and port")
|
|
||||||
flag.Parse()
|
|
||||||
|
|
||||||
// check and parse address
|
if err := svc.Run(prg); err != nil {
|
||||||
conn, err := net.ResolveTCPAddr("tcp", connect)
|
|
||||||
if err != nil {
|
|
||||||
flag.PrintDefaults()
|
|
||||||
log.Fatal(err)
|
log.Fatal(err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// listen on address
|
|
||||||
ln, err := net.Listen("tcp", listen)
|
|
||||||
if err != nil {
|
|
||||||
flag.PrintDefaults()
|
|
||||||
log.Fatal(err)
|
|
||||||
}
|
|
||||||
|
|
||||||
log.Printf("listening on %v", ln.Addr())
|
|
||||||
log.Printf("will connect to %v", conn)
|
|
||||||
|
|
||||||
for i := 0; ; i++ {
|
|
||||||
// accept new connection
|
|
||||||
c, err := ln.Accept()
|
|
||||||
if err != nil {
|
|
||||||
log.Fatal(err)
|
|
||||||
}
|
|
||||||
|
|
||||||
log.Printf("connection %v from %v", i, c.RemoteAddr())
|
|
||||||
|
|
||||||
cn, err := net.DialTCP("tcp", nil, conn)
|
|
||||||
if err != nil {
|
|
||||||
c.Close()
|
|
||||||
log.Print(err)
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
|
|
||||||
go pipe(c, cn, i)
|
|
||||||
go pipe(cn, c, i)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func pipe(w io.WriteCloser, r io.ReadCloser, count int) {
|
func (p *program) Init(env svc.Environment) error {
|
||||||
n, err := io.Copy(w, r)
|
if env.IsWindowsService() {
|
||||||
|
dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
logPath := filepath.Join(dir, "forward.log")
|
||||||
|
|
||||||
r.Close()
|
f, err := os.OpenFile(logPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
|
||||||
w.Close()
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
log.Printf("connection %v closed, %v bytes", count, n)
|
p.logFile = f
|
||||||
|
log.SetOutput(f)
|
||||||
opError, ok := err.(*net.OpError)
|
|
||||||
if err != nil && (!ok || opError.Op != "readfrom") {
|
|
||||||
log.Printf("warning! %v", err)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
setup()
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *program) Start() error {
|
||||||
|
p.wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
serve()
|
||||||
|
p.wg.Done()
|
||||||
|
}()
|
||||||
|
|
||||||
|
log.Print("started")
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p *program) Stop() error {
|
||||||
|
log.Print("stopping...")
|
||||||
|
|
||||||
|
ln.Close()
|
||||||
|
p.wg.Wait()
|
||||||
|
|
||||||
|
log.Print("stopped")
|
||||||
|
return nil
|
||||||
}
|
}
|
||||||
|
|
|
@ -0,0 +1,76 @@
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"flag"
|
||||||
|
"io"
|
||||||
|
"log"
|
||||||
|
"net"
|
||||||
|
)
|
||||||
|
|
||||||
|
var listen string
|
||||||
|
var connect string
|
||||||
|
|
||||||
|
var ln net.Listener
|
||||||
|
var conn *net.TCPAddr
|
||||||
|
|
||||||
|
func setup() {
|
||||||
|
flag.StringVar(&listen, "listen", ":8000", "listen on ip and port")
|
||||||
|
flag.StringVar(&connect, "connect", "", "forward to ip and port")
|
||||||
|
flag.Parse()
|
||||||
|
|
||||||
|
var err error
|
||||||
|
|
||||||
|
// check and parse address
|
||||||
|
conn, err = net.ResolveTCPAddr("tcp", connect)
|
||||||
|
if err != nil {
|
||||||
|
flag.PrintDefaults()
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// listen on address
|
||||||
|
ln, err = net.Listen("tcp", listen)
|
||||||
|
if err != nil {
|
||||||
|
flag.PrintDefaults()
|
||||||
|
log.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
log.Printf("listening on %v", ln.Addr())
|
||||||
|
log.Printf("will connect to %v", conn)
|
||||||
|
}
|
||||||
|
|
||||||
|
func serve() {
|
||||||
|
for i := 0; ; i++ {
|
||||||
|
// accept new connection
|
||||||
|
c, err := ln.Accept()
|
||||||
|
if err != nil {
|
||||||
|
log.Print(err)
|
||||||
|
break
|
||||||
|
}
|
||||||
|
|
||||||
|
log.Printf("connection %v from %v", i, c.RemoteAddr())
|
||||||
|
|
||||||
|
cn, err := net.DialTCP("tcp", nil, conn)
|
||||||
|
if err != nil {
|
||||||
|
c.Close()
|
||||||
|
log.Print(err)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
go pipe(c, cn, i)
|
||||||
|
go pipe(cn, c, i)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func pipe(w io.WriteCloser, r io.ReadCloser, count int) {
|
||||||
|
n, err := io.Copy(w, r)
|
||||||
|
|
||||||
|
r.Close()
|
||||||
|
w.Close()
|
||||||
|
|
||||||
|
log.Printf("connection %v closed, %v bytes", count, n)
|
||||||
|
|
||||||
|
opError, ok := err.(*net.OpError)
|
||||||
|
if err != nil && (!ok || opError.Op != "readfrom") {
|
||||||
|
log.Printf("warning! %v", err)
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in New Issue