Browse Source

Support Windows service with go-svc

master
Ambrose Chua 3 years ago
parent
commit
f6449feb55
  1. 6
      go.mod
  2. 84
      main.go
  3. 76
      server.go

6
go.mod

@ -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
)

84
main.go

@ -1,68 +1,66 @@
package main
import (
"flag"
"io"
"log"
"net"
"os"
"path/filepath"
"sync"
"github.com/judwhite/go-svc/svc"
)
var listen string
var connect string
type program struct {
logFile *os.File
wg sync.WaitGroup
}
func main() {
flag.StringVar(&listen, "listen", ":8000", "listen on ip and port")
flag.StringVar(&connect, "connect", "", "forward to ip and port")
flag.Parse()
prg := &program{}
// check and parse address
conn, err := net.ResolveTCPAddr("tcp", connect)
if err != nil {
flag.PrintDefaults()
if err := svc.Run(prg); err != nil {
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()
func (p *program) Init(env svc.Environment) error {
if env.IsWindowsService() {
dir, err := filepath.Abs(filepath.Dir(os.Args[0]))
if err != nil {
log.Fatal(err)
return err
}
logPath := filepath.Join(dir, "forward.log")
log.Printf("connection %v from %v", i, c.RemoteAddr())
cn, err := net.DialTCP("tcp", nil, conn)
f, err := os.OpenFile(logPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
if err != nil {
c.Close()
log.Print(err)
continue
return err
}
go pipe(c, cn, i)
go pipe(cn, c, i)
p.logFile = f
log.SetOutput(f)
}
setup()
return nil
}
func pipe(w io.WriteCloser, r io.ReadCloser, count int) {
n, err := io.Copy(w, r)
func (p *program) Start() error {
p.wg.Add(1)
go func() {
serve()
p.wg.Done()
}()
r.Close()
w.Close()
log.Print("started")
return nil
}
log.Printf("connection %v closed, %v bytes", count, n)
func (p *program) Stop() error {
log.Print("stopping...")
opError, ok := err.(*net.OpError)
if err != nil && (!ok || opError.Op != "readfrom") {
log.Printf("warning! %v", err)
}
ln.Close()
p.wg.Wait()
log.Print("stopped")
return nil
}

76
server.go

@ -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…
Cancel
Save