mirror of
https://github.com/fatedier/frp.git
synced 2026-01-11 22:23:12 +00:00
start refactoring
This commit is contained in:
fatedier
38
models/proto/tcp/process.go
Normal file
38
models/proto/tcp/process.go
Normal file
@@ -0,0 +1,38 @@
|
||||
// Copyright 2016 fatedier, fatedier@gmail.com
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package tcp
|
||||
|
||||
import (
|
||||
"io"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Join two io.ReadWriteCloser and do some operations.
|
||||
func Join(c1 io.ReadWriteCloser, c2 io.ReadWriteCloser) (inCount int64, outCount int64) {
|
||||
var wait sync.WaitGroup
|
||||
pipe := func(to io.ReadWriteCloser, from io.ReadWriteCloser, count *int64) {
|
||||
defer to.Close()
|
||||
defer from.Close()
|
||||
defer wait.Done()
|
||||
|
||||
*count, _ = io.Copy(to, from)
|
||||
}
|
||||
|
||||
wait.Add(2)
|
||||
go pipe(c1, c2, &inCount)
|
||||
go pipe(c2, c1, &outCount)
|
||||
wait.Wait()
|
||||
return
|
||||
}
|
||||
129
models/proto/tcp/process_test.go
Normal file
129
models/proto/tcp/process_test.go
Normal file
@@ -0,0 +1,129 @@
|
||||
// Copyright 2017 fatedier, fatedier@gmail.com
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package tcp
|
||||
|
||||
import (
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
|
||||
"github.com/fatedier/frp/utils/crypto"
|
||||
)
|
||||
|
||||
func TestJoin(t *testing.T) {
|
||||
assert := assert.New(t)
|
||||
|
||||
var (
|
||||
n int
|
||||
err error
|
||||
)
|
||||
text1 := "A document that gives tips for writing clear, idiomatic Go code. A must read for any new Go programmer. It augments the tour and the language specification, both of which should be read first."
|
||||
text2 := "A document that specifies the conditions under which reads of a variable in one goroutine can be guaranteed to observe values produced by writes to the same variable in a different goroutine."
|
||||
|
||||
// Forward bytes directly.
|
||||
pr, pw := io.Pipe()
|
||||
pr2, pw2 := io.Pipe()
|
||||
pr3, pw3 := io.Pipe()
|
||||
pr4, pw4 := io.Pipe()
|
||||
|
||||
conn1 := WrapReadWriteCloser(pr, pw2)
|
||||
conn2 := WrapReadWriteCloser(pr2, pw)
|
||||
conn3 := WrapReadWriteCloser(pr3, pw4)
|
||||
conn4 := WrapReadWriteCloser(pr4, pw3)
|
||||
|
||||
go func() {
|
||||
Join(conn2, conn3)
|
||||
}()
|
||||
|
||||
buf1 := make([]byte, 1024)
|
||||
buf2 := make([]byte, 1024)
|
||||
|
||||
conn1.Write([]byte(text1))
|
||||
conn4.Write([]byte(text2))
|
||||
|
||||
n, err = conn4.Read(buf1)
|
||||
assert.NoError(err)
|
||||
assert.Equal(text1, string(buf1[:n]))
|
||||
|
||||
n, err = conn1.Read(buf2)
|
||||
assert.NoError(err)
|
||||
assert.Equal(text2, string(buf2[:n]))
|
||||
|
||||
conn1.Close()
|
||||
conn2.Close()
|
||||
conn3.Close()
|
||||
conn4.Close()
|
||||
}
|
||||
|
||||
func TestJoinEncrypt(t *testing.T) {
|
||||
assert := assert.New(t)
|
||||
|
||||
var (
|
||||
n int
|
||||
err error
|
||||
)
|
||||
text1 := "1234567890"
|
||||
text2 := "abcdefghij"
|
||||
key := "authkey"
|
||||
|
||||
// Forward enrypted bytes.
|
||||
pr, pw := io.Pipe()
|
||||
pr2, pw2 := io.Pipe()
|
||||
pr3, pw3 := io.Pipe()
|
||||
pr4, pw4 := io.Pipe()
|
||||
pr5, pw5 := io.Pipe()
|
||||
pr6, pw6 := io.Pipe()
|
||||
|
||||
conn1 := WrapReadWriteCloser(pr, pw2)
|
||||
conn2 := WrapReadWriteCloser(pr2, pw)
|
||||
conn3 := WrapReadWriteCloser(pr3, pw4)
|
||||
conn4 := WrapReadWriteCloser(pr4, pw3)
|
||||
conn5 := WrapReadWriteCloser(pr5, pw6)
|
||||
conn6 := WrapReadWriteCloser(pr6, pw5)
|
||||
|
||||
r1, err := crypto.NewReader(conn3, []byte(key))
|
||||
assert.NoError(err)
|
||||
w1, err := crypto.NewWriter(conn3, []byte(key))
|
||||
assert.NoError(err)
|
||||
|
||||
r2, err := crypto.NewReader(conn4, []byte(key))
|
||||
assert.NoError(err)
|
||||
w2, err := crypto.NewWriter(conn4, []byte(key))
|
||||
assert.NoError(err)
|
||||
|
||||
go Join(conn2, WrapReadWriteCloser(r1, w1))
|
||||
go Join(WrapReadWriteCloser(r2, w2), conn5)
|
||||
|
||||
buf := make([]byte, 128)
|
||||
|
||||
conn1.Write([]byte(text1))
|
||||
conn6.Write([]byte(text2))
|
||||
|
||||
n, err = conn6.Read(buf)
|
||||
assert.NoError(err)
|
||||
assert.Equal(text1, string(buf[:n]))
|
||||
|
||||
n, err = conn1.Read(buf)
|
||||
assert.NoError(err)
|
||||
assert.Equal(text2, string(buf[:n]))
|
||||
|
||||
conn1.Close()
|
||||
conn2.Close()
|
||||
conn3.Close()
|
||||
conn4.Close()
|
||||
conn5.Close()
|
||||
conn6.Close()
|
||||
}
|
||||
89
models/proto/tcp/tcp.go
Normal file
89
models/proto/tcp/tcp.go
Normal file
@@ -0,0 +1,89 @@
|
||||
// Copyright 2017 fatedier, fatedier@gmail.com
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package tcp
|
||||
|
||||
import (
|
||||
"io"
|
||||
|
||||
"github.com/golang/snappy"
|
||||
|
||||
"github.com/fatedier/frp/utils/crypto"
|
||||
)
|
||||
|
||||
func WithEncryption(rwc io.ReadWriteCloser, key []byte) (res io.ReadWriteCloser, err error) {
|
||||
var (
|
||||
r io.Reader
|
||||
w io.Writer
|
||||
)
|
||||
r, err = crypto.NewReader(rwc, key)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
w, err = crypto.NewWriter(rwc, key)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
res = WrapReadWriteCloser(r, w)
|
||||
return
|
||||
}
|
||||
|
||||
func WithCompression(rwc io.ReadWriteCloser) (res io.ReadWriteCloser) {
|
||||
var (
|
||||
r io.Reader
|
||||
w io.Writer
|
||||
)
|
||||
r = snappy.NewReader(rwc)
|
||||
w = snappy.NewWriter(rwc)
|
||||
res = WrapReadWriteCloser(r, w)
|
||||
return
|
||||
}
|
||||
|
||||
func WrapReadWriteCloser(r io.Reader, w io.Writer) io.ReadWriteCloser {
|
||||
return &ReadWriteCloser{
|
||||
r: r,
|
||||
w: w,
|
||||
}
|
||||
}
|
||||
|
||||
type ReadWriteCloser struct {
|
||||
r io.Reader
|
||||
w io.Writer
|
||||
}
|
||||
|
||||
func (rwc *ReadWriteCloser) Read(p []byte) (n int, err error) {
|
||||
return rwc.r.Read(p)
|
||||
}
|
||||
|
||||
func (rwc *ReadWriteCloser) Write(p []byte) (n int, err error) {
|
||||
return rwc.w.Write(p)
|
||||
}
|
||||
|
||||
func (rwc *ReadWriteCloser) Close() (errRet error) {
|
||||
var err error
|
||||
if rc, ok := rwc.r.(io.Closer); ok {
|
||||
err = rc.Close()
|
||||
if err != nil {
|
||||
errRet = err
|
||||
}
|
||||
}
|
||||
|
||||
if wc, ok := rwc.w.(io.Closer); ok {
|
||||
err = wc.Close()
|
||||
if err != nil {
|
||||
errRet = err
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
100
models/proto/tcp/tcp_test.go
Normal file
100
models/proto/tcp/tcp_test.go
Normal file
@@ -0,0 +1,100 @@
|
||||
// Copyright 2017 fatedier, fatedier@gmail.com
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package tcp
|
||||
|
||||
import (
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestWithCompression(t *testing.T) {
|
||||
assert := assert.New(t)
|
||||
|
||||
// Forward compression bytes.
|
||||
pr, pw := io.Pipe()
|
||||
pr2, pw2 := io.Pipe()
|
||||
|
||||
conn1 := WrapReadWriteCloser(pr, pw2)
|
||||
conn2 := WrapReadWriteCloser(pr2, pw)
|
||||
|
||||
compressionStream1 := WithCompression(conn1)
|
||||
compressionStream2 := WithCompression(conn2)
|
||||
|
||||
var (
|
||||
n int
|
||||
err error
|
||||
)
|
||||
|
||||
text := "1234567812345678"
|
||||
buf := make([]byte, 256)
|
||||
|
||||
go compressionStream1.Write([]byte(text))
|
||||
n, err = compressionStream2.Read(buf)
|
||||
assert.NoError(err)
|
||||
assert.Equal(text, string(buf[:n]))
|
||||
|
||||
go compressionStream2.Write([]byte(text))
|
||||
n, err = compressionStream1.Read(buf)
|
||||
assert.NoError(err)
|
||||
assert.Equal(text, string(buf[:n]))
|
||||
}
|
||||
|
||||
func TestWithEncryption(t *testing.T) {
|
||||
assert := assert.New(t)
|
||||
var (
|
||||
n int
|
||||
err error
|
||||
)
|
||||
text1 := "Go is expressive, concise, clean, and efficient. Its concurrency mechanisms make it easy to write programs that get the most out of multicore and networked machines, while its novel type system enables flexible and modular program construction. Go compiles quickly to machine code yet has the convenience of garbage collection and the power of run-time reflection. It's a fast, statically typed, compiled language that feels like a dynamically typed, interpreted language."
|
||||
text2 := "An interactive introduction to Go in three sections. The first section covers basic syntax and data structures; the second discusses methods and interfaces; and the third introduces Go's concurrency primitives. Each section concludes with a few exercises so you can practice what you've learned. You can take the tour online or install it locally with"
|
||||
key := "authkey"
|
||||
|
||||
// Forward enrypted bytes.
|
||||
pr, pw := io.Pipe()
|
||||
pr2, pw2 := io.Pipe()
|
||||
pr3, pw3 := io.Pipe()
|
||||
pr4, pw4 := io.Pipe()
|
||||
pr5, pw5 := io.Pipe()
|
||||
pr6, pw6 := io.Pipe()
|
||||
|
||||
conn1 := WrapReadWriteCloser(pr, pw2)
|
||||
conn2 := WrapReadWriteCloser(pr2, pw)
|
||||
conn3 := WrapReadWriteCloser(pr3, pw4)
|
||||
conn4 := WrapReadWriteCloser(pr4, pw3)
|
||||
conn5 := WrapReadWriteCloser(pr5, pw6)
|
||||
conn6 := WrapReadWriteCloser(pr6, pw5)
|
||||
|
||||
encryptStream1, err := WithEncryption(conn3, []byte(key))
|
||||
assert.NoError(err)
|
||||
encryptStream2, err := WithEncryption(conn4, []byte(key))
|
||||
assert.NoError(err)
|
||||
|
||||
go Join(conn2, encryptStream1)
|
||||
go Join(encryptStream2, conn5)
|
||||
|
||||
buf := make([]byte, 1024)
|
||||
|
||||
conn1.Write([]byte(text1))
|
||||
conn6.Write([]byte(text2))
|
||||
|
||||
n, err = conn6.Read(buf)
|
||||
assert.NoError(err)
|
||||
assert.Equal(text1, string(buf[:n]))
|
||||
|
||||
n, err = conn1.Read(buf)
|
||||
assert.NoError(err)
|
||||
}
|
||||
72
models/proto/udp/udp.go
Normal file
72
models/proto/udp/udp.go
Normal file
@@ -0,0 +1,72 @@
|
||||
// Copyright 2016 fatedier, fatedier@gmail.com
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package udp
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"net"
|
||||
)
|
||||
|
||||
type UdpPacket struct {
|
||||
Content []byte `json:"-"`
|
||||
Src *net.UDPAddr `json:"-"`
|
||||
Dst *net.UDPAddr `json:"-"`
|
||||
|
||||
EncodeContent string `json:"content"`
|
||||
SrcStr string `json:"src"`
|
||||
DstStr string `json:"dst"`
|
||||
}
|
||||
|
||||
func NewUdpPacket(content []byte, src, dst *net.UDPAddr) *UdpPacket {
|
||||
up := &UdpPacket{
|
||||
Src: src,
|
||||
Dst: dst,
|
||||
EncodeContent: base64.StdEncoding.EncodeToString(content),
|
||||
SrcStr: src.String(),
|
||||
DstStr: dst.String(),
|
||||
}
|
||||
return up
|
||||
}
|
||||
|
||||
// parse one udp packet struct to bytes
|
||||
func (up *UdpPacket) Pack() []byte {
|
||||
b, _ := json.Marshal(up)
|
||||
return b
|
||||
}
|
||||
|
||||
// parse from bytes to UdpPacket struct
|
||||
func (up *UdpPacket) UnPack(packet []byte) error {
|
||||
err := json.Unmarshal(packet, &up)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
up.Content, err = base64.StdEncoding.DecodeString(up.EncodeContent)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
up.Src, err = net.ResolveUDPAddr("udp", up.SrcStr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
up.Dst, err = net.ResolveUDPAddr("udp", up.DstStr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
50
models/proto/udp/udp_test.go
Normal file
50
models/proto/udp/udp_test.go
Normal file
@@ -0,0 +1,50 @@
|
||||
// Copyright 2016 fatedier, fatedier@gmail.com
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package udp
|
||||
|
||||
import (
|
||||
"net"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
var (
|
||||
content string = "udp packet test"
|
||||
src string = "1.1.1.1:1000"
|
||||
dst string = "2.2.2.2:2000"
|
||||
|
||||
udpMsg *UdpPacket
|
||||
)
|
||||
|
||||
func init() {
|
||||
srcAddr, _ := net.ResolveUDPAddr("udp", src)
|
||||
dstAddr, _ := net.ResolveUDPAddr("udp", dst)
|
||||
udpMsg = NewUdpPacket([]byte(content), srcAddr, dstAddr)
|
||||
}
|
||||
|
||||
func TestPack(t *testing.T) {
|
||||
assert := assert.New(t)
|
||||
msg := udpMsg.Pack()
|
||||
assert.Equal(string(msg), `{"content":"dWRwIHBhY2tldCB0ZXN0","src":"1.1.1.1:1000","dst":"2.2.2.2:2000"}`)
|
||||
}
|
||||
|
||||
func TestUnpack(t *testing.T) {
|
||||
assert := assert.New(t)
|
||||
udpMsg.UnPack([]byte(`{"content":"dWRwIHBhY2tldCB0ZXN0","src":"1.1.1.1:1000","dst":"2.2.2.2:2000"}`))
|
||||
assert.Equal(content, string(udpMsg.Content))
|
||||
assert.Equal(src, udpMsg.Src.String())
|
||||
assert.Equal(dst, udpMsg.Dst.String())
|
||||
}
|
||||
Reference in New Issue
Block a user