Skip to content

Commit

Permalink
Merge pull request #2 from systemli/Implement-Domain-Lookup
Browse files Browse the repository at this point in the history
✨ Implement Domain Lookup
  • Loading branch information
0x46616c6b authored Oct 2, 2024
2 parents 3e83dbf + 1cac90a commit 1c6a06e
Showing 10 changed files with 342 additions and 37 deletions.
32 changes: 4 additions & 28 deletions alias.go
Original file line number Diff line number Diff line change
@@ -4,42 +4,18 @@ import (
"bytes"
"fmt"
"net"
"os"
"strings"
)

type Alias struct {
address string
userli UserliService
userli UserliService
}

func NewAlias(address string, userli UserliService) *Alias {
return &Alias{address: address, userli: userli}
func NewAlias(userli UserliService) *Alias {
return &Alias{userli: userli}
}

func (a *Alias) Listen() {
listen, err := net.Listen("tcp", a.address)
if err != nil {
fmt.Println("Error listening:", err.Error())
os.Exit(1)
}

defer listen.Close()
fmt.Println("Alias service listening on port", a.address)

for {
conn, err := listen.Accept()
if err != nil {
fmt.Println("Error accepting: ", err.Error())
os.Exit(1)
}

go a.handle(conn)
}
}

func (a *Alias) handle(conn net.Conn) {
defer conn.Close()
func (a *Alias) Handle(conn net.Conn) {
command := make([]byte, 4096)
_, err := conn.Read(command)
if err != nil {
19 changes: 17 additions & 2 deletions alias_test.go
Original file line number Diff line number Diff line change
@@ -2,17 +2,31 @@ package main

import (
"bytes"
"context"
"crypto/rand"
"errors"
"math/big"
"net"
"sync"
"testing"

"github.com/stretchr/testify/suite"
)

type AliasTestSuite struct {
suite.Suite

wg *sync.WaitGroup
ctx context.Context
}

func (s *AliasTestSuite) SetupTest() {
s.wg = &sync.WaitGroup{}
s.ctx = context.Background()
}

func (s *AliasTestSuite) AfterTest(_, _ string) {
s.ctx.Done()
}

func (s *AliasTestSuite) TestAlias() {
@@ -24,8 +38,9 @@ func (s *AliasTestSuite) TestAlias() {
portNumber, _ := rand.Int(rand.Reader, big.NewInt(65535-20000))
listen := ":" + portNumber.String()

alias := NewAlias(listen, userli)
go alias.Listen()
alias := NewAlias(userli)

go StartTCPServer(s.ctx, s.wg, listen, alias.Handle)

// wait until the server is ready
for {
48 changes: 48 additions & 0 deletions domain.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
package main

import (
"bytes"
"fmt"
"net"
"strings"
)

type Domain struct {
userli UserliService
}

func NewDomain(userli UserliService) *Domain {
return &Domain{userli: userli}
}

func (d *Domain) Handle(conn net.Conn) {
defer conn.Close()
command := make([]byte, 4096)
_, err := conn.Read(command)
if err != nil {
fmt.Println("Error reading:", err.Error())
}

command = bytes.Trim(command, "\x00")
parts := strings.Split(string(command), " ")
if len(parts) < 2 || parts[0] != "get" {
fmt.Println("Invalid command")
_, _ = conn.Write([]byte("400 Bad Request\n"))
return
}

domain := strings.TrimSuffix(parts[1], "\n")
exists, err := d.userli.GetDomain(string(domain))
if err != nil {
fmt.Println("Error fetching domain:", err.Error())
_, _ = conn.Write([]byte("400 Error fetching domain\n"))
return
}

if !exists {
_, _ = conn.Write([]byte("500 NO%20RESULT\n"))
return
}

_, _ = conn.Write([]byte("200 1\n"))
}
105 changes: 105 additions & 0 deletions domain_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
package main

import (
"bytes"
"context"
"crypto/rand"
"errors"
"math/big"
"net"
"sync"
"testing"

"github.com/stretchr/testify/suite"
)

type DomainTestSuite struct {
suite.Suite

wg *sync.WaitGroup
ctx context.Context
}

func (s *DomainTestSuite) SetupTest() {
s.wg = &sync.WaitGroup{}
s.ctx = context.Background()
}

func (s *DomainTestSuite) AfterTest(_, _ string) {
s.ctx.Done()
}

func (s *DomainTestSuite) TestDomain() {
userli := new(MockUserliService)
userli.On("GetDomain", "example.com").Return(true, nil)
userli.On("GetDomain", "notfound.com").Return(false, nil)
userli.On("GetDomain", "error.com").Return(false, errors.New("error"))

portNumber, _ := rand.Int(rand.Reader, big.NewInt(65535-20000))
listen := ":" + portNumber.String()

domain := NewDomain(userli)

go StartTCPServer(s.ctx, s.wg, listen, domain.Handle)

// wait until the server is ready
for {
conn, err := net.Dial("tcp", listen)
if err == nil {
conn.Close()
break
}
}

s.Run("success", func() {
conn, err := net.Dial("tcp", listen)
s.NoError(err)

_, err = conn.Write([]byte("get example.com"))
s.NoError(err)

response := make([]byte, 4096)
_, err = conn.Read(response)
s.NoError(err)

s.Equal("200 1\n", string(bytes.Trim(response, "\x00")))

conn.Close()
})

s.Run("not found", func() {
conn, err := net.Dial("tcp", listen)
s.NoError(err)

_, err = conn.Write([]byte("get notfound.com"))
s.NoError(err)

response := make([]byte, 4096)
_, err = conn.Read(response)
s.NoError(err)

s.Equal("500 NO%20RESULT\n", string(bytes.Trim(response, "\x00")))

conn.Close()
})

s.Run("error", func() {
conn, err := net.Dial("tcp", listen)
s.NoError(err)

_, err = conn.Write([]byte("get error.com"))
s.NoError(err)

response := make([]byte, 4096)
_, err = conn.Read(response)
s.NoError(err)

s.Equal("400 Error fetching domain\n", string(bytes.Trim(response, "\x00")))

conn.Close()
})
}

func TestDomain(t *testing.T) {
suite.Run(t, new(DomainTestSuite))
}
10 changes: 10 additions & 0 deletions handler.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
package main

import "net"

// PostfixHandler is an interface that defines the Handle method
// for handling postfix postmap requests.
// See https://www.postfix.org/postmap.1.html
type PostfixHandler interface {
Handle(conn net.Conn)
}
24 changes: 22 additions & 2 deletions main.go
Original file line number Diff line number Diff line change
@@ -1,8 +1,12 @@
package main

import (
"context"
"fmt"
"os"
"os/signal"
"sync"
"syscall"
)

var (
@@ -28,8 +32,24 @@ func main() {
aliasListenAddr = ":10001"
}

domainListenAddr := os.Getenv("DOMAIN_LISTEN_ADDR")
if domainListenAddr == "" {
domainListenAddr = ":10002"
}

userli := NewUserli(userliToken, userliBaseURL)
alias := NewAlias(aliasListenAddr, userli)
alias := NewAlias(userli)
domain := NewDomain(userli)

ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer stop()

var wg sync.WaitGroup

wg.Add(2)
go StartTCPServer(ctx, &wg, aliasListenAddr, alias.Handle)
go StartTCPServer(ctx, &wg, domainListenAddr, domain.Handle)

alias.Listen()
wg.Wait()
fmt.Println("Servers stopped")
}
30 changes: 29 additions & 1 deletion mock_UserliService.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

38 changes: 38 additions & 0 deletions server.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
package main

import (
"context"
"fmt"
"net"
"sync"
)

func StartTCPServer(ctx context.Context, wg *sync.WaitGroup, addr string, handler func(net.Conn)) {
defer wg.Done()

listener, err := net.Listen("tcp", addr)
if err != nil {
fmt.Println("Error creating listener:", err.Error())
return
}
defer listener.Close()

go func() {
<-ctx.Done()
listener.Close()
}()

for {
conn, err := listener.Accept()
if err != nil {
if ctx.Err() != nil {
fmt.Println("Server stopped on port ", addr)
return
}
fmt.Println("Error accepting connection:", err.Error())
continue
}

go handler(conn)
}
}
Loading

0 comments on commit 1c6a06e

Please sign in to comment.