mirror of https://github.com/fatedier/frp.git
vendor: add package golib/net
fatedier
7 years ago
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