mirror of https://github.com/fatedier/frp.git
vendor: add package golib/net
parent
a27e3dda88
commit
20fcb58437
@ -1,95 +0,0 @@
|
|||||||
package mux
|
|
||||||
|
|
||||||
import (
|
|
||||||
"bufio"
|
|
||||||
"io/ioutil"
|
|
||||||
"net"
|
|
||||||
"net/http"
|
|
||||||
"net/http/httptest"
|
|
||||||
"testing"
|
|
||||||
"time"
|
|
||||||
|
|
||||||
"github.com/stretchr/testify/assert"
|
|
||||||
)
|
|
||||||
|
|
||||||
func runHttpSvr(ln net.Listener) *httptest.Server {
|
|
||||||
svr := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
||||||
w.Write([]byte("http service"))
|
|
||||||
}))
|
|
||||||
svr.Listener = ln
|
|
||||||
svr.Start()
|
|
||||||
return svr
|
|
||||||
}
|
|
||||||
|
|
||||||
func runHttpsSvr(ln net.Listener) *httptest.Server {
|
|
||||||
svr := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
||||||
w.Write([]byte("https service"))
|
|
||||||
}))
|
|
||||||
svr.Listener = ln
|
|
||||||
svr.StartTLS()
|
|
||||||
return svr
|
|
||||||
}
|
|
||||||
|
|
||||||
func runEchoSvr(ln net.Listener) {
|
|
||||||
go func() {
|
|
||||||
for {
|
|
||||||
conn, err := ln.Accept()
|
|
||||||
if err != nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
rd := bufio.NewReader(conn)
|
|
||||||
data, err := rd.ReadString('\n')
|
|
||||||
if err != nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
conn.Write([]byte(data))
|
|
||||||
conn.Close()
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
}
|
|
||||||
|
|
||||||
func TestMux(t *testing.T) {
|
|
||||||
assert := assert.New(t)
|
|
||||||
|
|
||||||
ln, err := net.Listen("tcp", "127.0.0.1:")
|
|
||||||
assert.NoError(err)
|
|
||||||
|
|
||||||
mux := NewMux()
|
|
||||||
httpLn := mux.ListenHttp(0)
|
|
||||||
httpsLn := mux.ListenHttps(0)
|
|
||||||
defaultLn := mux.DefaultListener()
|
|
||||||
go mux.Serve(ln)
|
|
||||||
time.Sleep(100 * time.Millisecond)
|
|
||||||
|
|
||||||
httpSvr := runHttpSvr(httpLn)
|
|
||||||
defer httpSvr.Close()
|
|
||||||
httpsSvr := runHttpsSvr(httpsLn)
|
|
||||||
defer httpsSvr.Close()
|
|
||||||
runEchoSvr(defaultLn)
|
|
||||||
defer ln.Close()
|
|
||||||
|
|
||||||
// test http service
|
|
||||||
resp, err := http.Get(httpSvr.URL)
|
|
||||||
assert.NoError(err)
|
|
||||||
data, err := ioutil.ReadAll(resp.Body)
|
|
||||||
assert.NoError(err)
|
|
||||||
assert.Equal("http service", string(data))
|
|
||||||
|
|
||||||
// test https service
|
|
||||||
client := httpsSvr.Client()
|
|
||||||
resp, err = client.Get(httpsSvr.URL)
|
|
||||||
assert.NoError(err)
|
|
||||||
data, err = ioutil.ReadAll(resp.Body)
|
|
||||||
assert.NoError(err)
|
|
||||||
assert.Equal("https service", string(data))
|
|
||||||
|
|
||||||
// test echo service
|
|
||||||
conn, err := net.Dial("tcp", ln.Addr().String())
|
|
||||||
assert.NoError(err)
|
|
||||||
_, err = conn.Write([]byte("test echo\n"))
|
|
||||||
assert.NoError(err)
|
|
||||||
data = make([]byte, 1024)
|
|
||||||
n, err := conn.Read(data)
|
|
||||||
assert.NoError(err)
|
|
||||||
assert.Equal("test echo\n", string(data[:n]))
|
|
||||||
}
|
|
@ -0,0 +1,64 @@
|
|||||||
|
// Copyright 2018 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 net
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
)
|
||||||
|
|
||||||
|
type SharedConn struct {
|
||||||
|
net.Conn
|
||||||
|
buf *bytes.Buffer
|
||||||
|
}
|
||||||
|
|
||||||
|
// the bytes you read in io.Reader, will be reserved in SharedConn
|
||||||
|
func NewSharedConn(conn net.Conn) (*SharedConn, io.Reader) {
|
||||||
|
sc := &SharedConn{
|
||||||
|
Conn: conn,
|
||||||
|
buf: bytes.NewBuffer(make([]byte, 0, 1024)),
|
||||||
|
}
|
||||||
|
return sc, io.TeeReader(conn, sc.buf)
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewSharedConnSize(conn net.Conn, bufSize int) (*SharedConn, io.Reader) {
|
||||||
|
sc := &SharedConn{
|
||||||
|
Conn: conn,
|
||||||
|
buf: bytes.NewBuffer(make([]byte, 0, bufSize)),
|
||||||
|
}
|
||||||
|
return sc, io.TeeReader(conn, sc.buf)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Not thread safety.
|
||||||
|
func (sc *SharedConn) Read(p []byte) (n int, err error) {
|
||||||
|
if sc.buf == nil {
|
||||||
|
return sc.Conn.Read(p)
|
||||||
|
}
|
||||||
|
n, err = sc.buf.Read(p)
|
||||||
|
if err == io.EOF {
|
||||||
|
sc.buf = nil
|
||||||
|
var n2 int
|
||||||
|
n2, err = sc.Conn.Read(p[n:])
|
||||||
|
n += n2
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
func (sc *SharedConn) ResetBuf(buffer []byte) (err error) {
|
||||||
|
sc.buf.Reset()
|
||||||
|
_, err = sc.buf.Write(buffer)
|
||||||
|
return err
|
||||||
|
}
|
Loading…
Reference in New Issue