153 lines
3.2 KiB
Go
153 lines
3.2 KiB
Go
/*
|
|
* Copyright 2025 The Dragonfly Authors
|
|
*
|
|
* 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 redis
|
|
|
|
import (
|
|
"bufio"
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
|
|
logger "d7y.io/dragonfly/v2/internal/dflog"
|
|
)
|
|
|
|
type Proxy interface {
|
|
Serve() error
|
|
Stop()
|
|
}
|
|
|
|
type proxy struct {
|
|
from string
|
|
to string
|
|
done chan struct{}
|
|
}
|
|
|
|
// NewProxy creates a new proxy instance for redirecting traffic to redis.
|
|
func NewProxy(from string, to string) Proxy {
|
|
return &proxy{
|
|
from: from,
|
|
to: to,
|
|
done: make(chan struct{}),
|
|
}
|
|
}
|
|
|
|
// Serve starts the proxy server and listens for incoming connections.
|
|
func (p *proxy) Serve() error {
|
|
listener, err := net.Listen("tcp", p.from)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer listener.Close()
|
|
|
|
for {
|
|
select {
|
|
case <-p.done:
|
|
return nil
|
|
default:
|
|
conn, err := listener.Accept()
|
|
if err != nil {
|
|
logger.Errorf("error accepting conn: %v", err)
|
|
} else {
|
|
go p.handleConn(conn)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Stop stops the proxy server and closes all connections.
|
|
func (p *proxy) Stop() {
|
|
if p.done == nil {
|
|
return
|
|
}
|
|
close(p.done)
|
|
p.done = nil
|
|
}
|
|
|
|
// handleConn handles the incoming connection and establishes a connection to the remote host.
|
|
func (p *proxy) handleConn(conn net.Conn) {
|
|
defer conn.Close()
|
|
|
|
reader, isRedisProtocol := p.isRedisProtocol(conn)
|
|
if !isRedisProtocol {
|
|
logger.Errorf("not a redis protocol: %s", conn.RemoteAddr())
|
|
return
|
|
}
|
|
|
|
rConn, err := net.Dial("tcp", p.to)
|
|
if err != nil {
|
|
logger.Errorf("error dialing remote host: %v", err)
|
|
return
|
|
}
|
|
defer rConn.Close()
|
|
|
|
wg := &sync.WaitGroup{}
|
|
wg.Add(2)
|
|
go p.copy(rConn, conn, wg)
|
|
go p.copyReader(reader, rConn, wg)
|
|
wg.Wait()
|
|
}
|
|
|
|
// copy copies data from one connection to another.
|
|
func (p *proxy) copy(from, to net.Conn, wg *sync.WaitGroup) {
|
|
defer wg.Done()
|
|
select {
|
|
case <-p.done:
|
|
return
|
|
default:
|
|
if _, err := io.Copy(to, from); err != nil {
|
|
logger.Errorf("error copying from %s to %s: %v", from.RemoteAddr(), to.RemoteAddr(), err)
|
|
p.Stop()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// copyReader copies data from a reader to a connection.
|
|
func (p *proxy) copyReader(from io.Reader, to net.Conn, wg *sync.WaitGroup) {
|
|
defer wg.Done()
|
|
select {
|
|
case <-p.done:
|
|
return
|
|
default:
|
|
if _, err := io.Copy(to, from); err != nil {
|
|
logger.Errorf("error copying to %s: %v", to.RemoteAddr(), err)
|
|
p.Stop()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// isRedisProtocol checks if the connection uses the Redis protocol.
|
|
func (p *proxy) isRedisProtocol(conn net.Conn) (io.Reader, bool) {
|
|
reader := bufio.NewReader(conn)
|
|
firstByte, err := reader.Peek(1)
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
logger.Errorf("reading first byte from client failed: %s: %v", conn.RemoteAddr(), err)
|
|
}
|
|
|
|
return reader, false
|
|
}
|
|
|
|
switch firstByte[0] {
|
|
case '*', '+', '-', ':', '$':
|
|
return reader, true
|
|
}
|
|
|
|
return reader, false
|
|
}
|