Change rwc to rwc with addresses

This commit is contained in:
9seconds
2018-07-02 16:12:31 +03:00
parent 6fa5f31ca8
commit 09476cc467
15 changed files with 109 additions and 68 deletions
+13 -9
View File
@@ -4,20 +4,20 @@ import (
"bytes"
"crypto/aes"
"crypto/cipher"
"io"
"net"
"github.com/juju/errors"
)
type BlockCipherReadWriteCloser struct {
type BlockCipherReadWriteCloserWithAddr struct {
buf *bytes.Buffer
conn io.ReadWriteCloser
conn ReadWriteCloserWithAddr
encryptor cipher.BlockMode
decryptor cipher.BlockMode
}
func (c *BlockCipherReadWriteCloser) Read(p []byte) (int, error) {
func (c *BlockCipherReadWriteCloserWithAddr) Read(p []byte) (int, error) {
if c.buf.Len() > 0 {
return c.flush(p)
}
@@ -34,7 +34,7 @@ func (c *BlockCipherReadWriteCloser) Read(p []byte) (int, error) {
return c.flush(p)
}
func (c *BlockCipherReadWriteCloser) Write(p []byte) (int, error) {
func (c *BlockCipherReadWriteCloserWithAddr) Write(p []byte) (int, error) {
if len(p)%aes.BlockSize > 0 {
return 0, errors.Errorf("Incorrect block size %d", len(p))
}
@@ -50,12 +50,16 @@ func (c *BlockCipherReadWriteCloser) Write(p []byte) (int, error) {
return c.conn.Write(encrypted)
}
func (c *BlockCipherReadWriteCloser) Close() error {
func (c *BlockCipherReadWriteCloserWithAddr) Close() error {
defer putBuffer(c.buf)
return c.conn.Close()
}
func (c *BlockCipherReadWriteCloser) flush(p []byte) (int, error) {
func (c *BlockCipherReadWriteCloserWithAddr) Addr() *net.TCPAddr {
return c.conn.Addr()
}
func (c *BlockCipherReadWriteCloserWithAddr) flush(p []byte) (int, error) {
sizeToRead := len(p)
if c.buf.Len() < sizeToRead {
sizeToRead = c.buf.Len()
@@ -76,8 +80,8 @@ func (c *BlockCipherReadWriteCloser) flush(p []byte) (int, error) {
return sizeToRead, nil
}
func NewBlockCipherRWC(conn io.ReadWriteCloser, encryptor, decryptor cipher.BlockMode) io.ReadWriteCloser {
return &BlockCipherReadWriteCloser{
func NewBlockCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.BlockMode) ReadWriteCloserWithAddr {
return &BlockCipherReadWriteCloserWithAddr{
buf: getBuffer(),
conn: conn,
encryptor: encryptor,
+12 -8
View File
@@ -2,21 +2,21 @@ package wrappers
import (
"context"
"io"
"net"
"github.com/juju/errors"
)
// CtxReadWriteCloser wraps underlying connection and does management of the
// context and its cancel function.
type CtxReadWriteCloser struct {
type CtxReadWriteCloserWithAddr struct {
ctx context.Context
conn io.ReadWriteCloser
conn ReadWriteCloserWithAddr
cancel context.CancelFunc
}
// Read reads from connection
func (c *CtxReadWriteCloser) Read(p []byte) (int, error) {
func (c *CtxReadWriteCloserWithAddr) Read(p []byte) (int, error) {
select {
case <-c.ctx.Done():
return 0, errors.Annotate(c.ctx.Err(), "Read is failed because of closed context")
@@ -30,7 +30,7 @@ func (c *CtxReadWriteCloser) Read(p []byte) (int, error) {
}
// Write writes into connection.
func (c *CtxReadWriteCloser) Write(p []byte) (int, error) {
func (c *CtxReadWriteCloserWithAddr) Write(p []byte) (int, error) {
select {
case <-c.ctx.Done():
return 0, errors.Annotate(c.ctx.Err(), "Write is failed because of closed context")
@@ -44,14 +44,18 @@ func (c *CtxReadWriteCloser) Write(p []byte) (int, error) {
}
// Close closes underlying connection.
func (c *CtxReadWriteCloser) Close() error {
func (c *CtxReadWriteCloserWithAddr) Close() error {
return c.conn.Close()
}
func (c *CtxReadWriteCloserWithAddr) Addr() *net.TCPAddr {
return c.conn.Addr()
}
// NewCtxRWC returns ReadWriteCloser which respects given context,
// cancellation etc.
func NewCtxRWC(ctx context.Context, cancel context.CancelFunc, conn io.ReadWriteCloser) io.ReadWriteCloser {
return &CtxReadWriteCloser{
func NewCtxRWC(ctx context.Context, cancel context.CancelFunc, conn ReadWriteCloserWithAddr) ReadWriteCloserWithAddr {
return &CtxReadWriteCloserWithAddr{
conn: conn,
ctx: ctx,
cancel: cancel,
+12 -8
View File
@@ -1,44 +1,48 @@
package wrappers
import (
"io"
"net"
"go.uber.org/zap"
)
// LogReadWriteCloser adds additional logging for reading/writing. All
// logging is performed for debug mode only.
type LogReadWriteCloser struct {
conn io.ReadWriteCloser
type LogReadWriteCloserWithAddr struct {
conn ReadWriteCloserWithAddr
logger *zap.SugaredLogger
sockid string
name string
}
// Read reads from connection
func (l *LogReadWriteCloser) Read(p []byte) (n int, err error) {
func (l *LogReadWriteCloserWithAddr) Read(p []byte) (n int, err error) {
n, err = l.conn.Read(p)
l.logger.Debugw("Finish reading", "name", l.name, "socketid", l.sockid, "nbytes", n, "error", err)
return
}
// Write writes into connection.
func (l *LogReadWriteCloser) Write(p []byte) (n int, err error) {
func (l *LogReadWriteCloserWithAddr) Write(p []byte) (n int, err error) {
n, err = l.conn.Write(p)
l.logger.Debugw("Finish writing", "name", l.name, "socketid", l.sockid, "nbytes", n, "error", err)
return
}
// Close closes underlying connection.
func (l *LogReadWriteCloser) Close() error {
func (l *LogReadWriteCloserWithAddr) Close() error {
err := l.conn.Close()
l.logger.Debugw("Finish closing socket", "name", l.name, "socketid", l.sockid, "error", err)
return err
}
func (l *LogReadWriteCloserWithAddr) Addr() *net.TCPAddr {
return l.conn.Addr()
}
// NewLogRWC wraps ReadWriteCloser with logger calls.
func NewLogRWC(conn io.ReadWriteCloser, logger *zap.SugaredLogger, sockid string, name string) io.ReadWriteCloser {
return &LogReadWriteCloser{
func NewLogRWC(conn ReadWriteCloserWithAddr, logger *zap.SugaredLogger, sockid string, name string) ReadWriteCloserWithAddr {
return &LogReadWriteCloserWithAddr{
conn: conn,
logger: logger,
sockid: sockid,
+12
View File
@@ -0,0 +1,12 @@
package wrappers
import (
"io"
"net"
)
type ReadWriteCloserWithAddr interface {
io.ReadWriteCloser
Addr() *net.TCPAddr
}
+12 -8
View File
@@ -2,26 +2,26 @@ package wrappers
import (
"crypto/cipher"
"io"
"net"
)
// StreamCipherReadWriteCloser is a ReadWriteCloser which ciphers
// incoming and outgoing data with givem cipher.Stream instances.
type StreamCipherReadWriteCloser struct {
type StreamCipherReadWriteCloserWithAddr struct {
encryptor cipher.Stream
decryptor cipher.Stream
conn io.ReadWriteCloser
conn ReadWriteCloserWithAddr
}
// Read reads from connection
func (c *StreamCipherReadWriteCloser) Read(p []byte) (n int, err error) {
func (c *StreamCipherReadWriteCloserWithAddr) Read(p []byte) (n int, err error) {
n, err = c.conn.Read(p)
c.decryptor.XORKeyStream(p, p[:n])
return
}
// Write writes into connection.
func (c *StreamCipherReadWriteCloser) Write(p []byte) (int, error) {
func (c *StreamCipherReadWriteCloserWithAddr) Write(p []byte) (int, error) {
// This is to decrease an amount of allocations. Unfortunately, escape
// analysis in (at least Golang 1.10) is absolutely not perfect. For
// example, it understands that we want to have a slice locally, right?
@@ -39,14 +39,18 @@ func (c *StreamCipherReadWriteCloser) Write(p []byte) (int, error) {
}
// Close closes underlying connection.
func (c *StreamCipherReadWriteCloser) Close() error {
func (c *StreamCipherReadWriteCloserWithAddr) Close() error {
return c.conn.Close()
}
func (c *StreamCipherReadWriteCloserWithAddr) Addr() *net.TCPAddr {
return c.conn.Addr()
}
// NewStreamCipherRWC returns wrapper which transparently
// encrypts/decrypts traffic with obfuscated2 protocol.
func NewStreamCipherRWC(conn io.ReadWriteCloser, encryptor, decryptor cipher.Stream) io.ReadWriteCloser {
return &StreamCipherReadWriteCloser{
func NewStreamCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.Stream) ReadWriteCloserWithAddr {
return &StreamCipherReadWriteCloserWithAddr{
conn: conn,
encryptor: encryptor,
decryptor: decryptor,
+10 -7
View File
@@ -1,31 +1,34 @@
package wrappers
import (
"io"
"net"
"time"
"github.com/9seconds/mtg/config"
)
type TimeoutReadWriteCloser struct {
type TimeoutReadWriteCloserWithAddr struct {
conn net.Conn
}
func (t *TimeoutReadWriteCloser) Read(p []byte) (int, error) {
func (t *TimeoutReadWriteCloserWithAddr) Read(p []byte) (int, error) {
t.conn.SetReadDeadline(time.Now().Add(config.TimeoutRead))
return t.conn.Read(p)
}
func (t *TimeoutReadWriteCloser) Write(p []byte) (int, error) {
func (t *TimeoutReadWriteCloserWithAddr) Write(p []byte) (int, error) {
t.conn.SetWriteDeadline(time.Now().Add(config.TimeoutWrite))
return t.conn.Write(p)
}
func (t *TimeoutReadWriteCloser) Close() error {
func (t *TimeoutReadWriteCloserWithAddr) Close() error {
return t.conn.Close()
}
func NewTimeoutRWC(conn net.Conn) io.ReadWriteCloser {
return &TimeoutReadWriteCloser{conn}
func (t *TimeoutReadWriteCloserWithAddr) Addr() *net.TCPAddr {
return t.conn.RemoteAddr().(*net.TCPAddr)
}
func NewTimeoutRWC(conn net.Conn) ReadWriteCloserWithAddr {
return &TimeoutReadWriteCloserWithAddr{conn}
}
+12 -8
View File
@@ -1,37 +1,41 @@
package wrappers
import "io"
import "net"
// TrafficReadWriteCloser counts an amount of ingress/egress traffic by
// calling given callbacks.
type TrafficReadWriteCloser struct {
conn io.ReadWriteCloser
type TrafficReadWriteCloserWithAddr struct {
conn ReadWriteCloserWithAddr
readCallback func(int)
writeCallback func(int)
}
// Read reads from connection
func (t *TrafficReadWriteCloser) Read(p []byte) (n int, err error) {
func (t *TrafficReadWriteCloserWithAddr) Read(p []byte) (n int, err error) {
n, err = t.conn.Read(p)
t.readCallback(n)
return
}
// Write writes into connection.
func (t *TrafficReadWriteCloser) Write(p []byte) (n int, err error) {
func (t *TrafficReadWriteCloserWithAddr) Write(p []byte) (n int, err error) {
n, err = t.conn.Write(p)
t.writeCallback(n)
return
}
// Close closes underlying connection.
func (t *TrafficReadWriteCloser) Close() error {
func (t *TrafficReadWriteCloserWithAddr) Close() error {
return t.conn.Close()
}
func (t *TrafficReadWriteCloserWithAddr) Addr() *net.TCPAddr {
return t.conn.Addr()
}
// NewTrafficRWC wraps ReadWriteCloser to have read/write callbacks.
func NewTrafficRWC(conn io.ReadWriteCloser, readCallback, writeCallback func(int)) io.ReadWriteCloser {
return &TrafficReadWriteCloser{
func NewTrafficRWC(conn ReadWriteCloserWithAddr, readCallback, writeCallback func(int)) ReadWriteCloserWithAddr {
return &TrafficReadWriteCloserWithAddr{
conn: conn,
readCallback: readCallback,
writeCallback: writeCallback,