60 lines
985 B
Go
60 lines
985 B
Go
|
package forward
|
||
|
|
|
||
|
|
import (
|
||
|
|
"errors"
|
||
|
|
"io"
|
||
|
|
"net"
|
||
|
|
|
||
|
|
"git.lovezsh.com/lovezsh/nip/internal/pkg/encoding"
|
||
|
|
"git.lovezsh.com/lovezsh/nip/internal/proto"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Forwarder interface {
|
||
|
|
Forward(net.Conn, string) (err error)
|
||
|
|
}
|
||
|
|
|
||
|
|
type forwarder struct {
|
||
|
|
server string
|
||
|
|
token string
|
||
|
|
}
|
||
|
|
|
||
|
|
func New(server string, token string) Forwarder {
|
||
|
|
return &forwarder{server, token}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (f *forwarder) Forward(src net.Conn, target string) (err error) {
|
||
|
|
dst, err := net.Dial("tcp", f.server)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
defer dst.Close()
|
||
|
|
|
||
|
|
req := proto.Request{
|
||
|
|
Target: target,
|
||
|
|
Token: f.token,
|
||
|
|
}
|
||
|
|
if err = encoding.NewEncoder(dst).Encode(&req); err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
var resp proto.Response
|
||
|
|
if err = encoding.NewDecoder(dst).Decode(&resp); err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
if resp.Code != 0 {
|
||
|
|
return errors.New(resp.Message)
|
||
|
|
}
|
||
|
|
|
||
|
|
ch := make(chan error)
|
||
|
|
go func() {
|
||
|
|
_, err := io.Copy(dst, src)
|
||
|
|
ch <- err
|
||
|
|
}()
|
||
|
|
go func() {
|
||
|
|
_, err := io.Copy(src, dst)
|
||
|
|
ch <- err
|
||
|
|
}()
|
||
|
|
return <-ch
|
||
|
|
}
|