path: root/internal/ipc/server.go

package ipc

import (
	"context"
	"encoding/json"
	"io"
	"log"
	"net"
	"sync"

	"mitmux/internal/rules"
	"mitmux/internal/store"
)

// Repeater sends raw bytes to scheme://host exactly as given and
// records the exchange to history - implemented by *proxy.Server.
type Repeater interface {
	Repeat(ctx context.Context, scheme, host string, raw []byte) (*store.Entry, error)
}

// Intruder runs a Sniper attack over a §marked§ request template -
// implemented by *proxy.Server.
type Intruder interface {
	Intrude(ctx context.Context, scheme, host string, template []byte, payloads []string,
		onResult func(position int, payload string, entry *store.Entry, sendErr error) bool) error
}

// Hub fans out newly captured history entries to subscribed clients.
type Hub struct {
	mu   sync.Mutex
	subs map[chan store.Summary]struct{}
}

// NewHub creates an empty Hub.
func NewHub() *Hub {
	return &Hub{subs: make(map[chan store.Summary]struct{})}
}

// Broadcast notifies all current subscribers of e. Slow subscribers
// drop entries rather than blocking the proxy.
func (h *Hub) Broadcast(e store.Summary) {
	h.mu.Lock()
	defer h.mu.Unlock()
	for ch := range h.subs {
		select {
		case ch <- e:
		default:
		}
	}
}

func (h *Hub) subscribe() chan store.Summary {
	ch := make(chan store.Summary, 64)
	h.mu.Lock()
	h.subs[ch] = struct{}{}
	h.mu.Unlock()
	return ch
}

func (h *Hub) unsubscribe(ch chan store.Summary) {
	h.mu.Lock()
	delete(h.subs, ch)
	h.mu.Unlock()
	close(ch)
}

// Server serves the daemon side of the mitmux control protocol.
type Server struct {
	db        *store.Store
	hub       *Hub
	repeater  Repeater
	intruder  Intruder
	proxyAddr string
}

// NewServer creates a control-protocol Server backed by db, broadcasting
// through hub and sending Repeater/Intruder requests through rep.
// proxyAddr is reported back to clients via the "status" request.
func NewServer(db *store.Store, hub *Hub, rep Repeater, proxyAddr string) *Server {
	s := &Server{db: db, hub: hub, repeater: rep, proxyAddr: proxyAddr}
	if in, ok := rep.(Intruder); ok {
		s.intruder = in
	}
	return s
}

// Serve accepts connections on ln until it returns an error (e.g. the
// listener is closed).
func (s *Server) Serve(ln net.Listener) error {
	for {
		conn, err := ln.Accept()
		if err != nil {
			return err
		}
		go s.handleConn(conn)
	}
}

func (s *Server) handleConn(conn net.Conn) {
	defer conn.Close()
	dec := json.NewDecoder(conn)
	enc := json.NewEncoder(conn)

	for {
		var req Request
		if err := dec.Decode(&req); err != nil {
			return
		}

		switch req.Type {
		case "status":
			count, err := s.db.Count()
			if err != nil {
				enc.Encode(Response{Type: "error", Error: err.Error()})
				continue
			}
			enc.Encode(Response{Type: "status", Status: &StatusMsg{ProxyAddr: s.proxyAddr, HistoryCount: count}})

		case "list":
			var entries []store.Summary
			var err error
			if req.Query != "" {
				entries, err = s.db.Search(req.Query, req.Limit, req.BeforeID)
			} else {
				entries, err = s.db.List(req.Limit, req.BeforeID)
			}
			if err != nil {
				enc.Encode(Response{Type: "error", Error: err.Error()})
				continue
			}
			enc.Encode(Response{Type: "list", Entries: entries})

		case "get":
			e, err := s.db.Get(req.ID)
			if err != nil {
				enc.Encode(Response{Type: "error", Error: err.Error()})
				continue
			}
			enc.Encode(Response{Type: "get", Detail: detailFromEntry(e)})

		case "repeat":
			if s.repeater == nil {
				enc.Encode(Response{Type: "error", Error: "repeater not available"})
				continue
			}
			e, err := s.repeater.Repeat(context.Background(), req.Scheme, req.Host, req.Raw)
			if err != nil {
				enc.Encode(Response{Type: "error", Error: err.Error()})
				continue
			}
			enc.Encode(Response{Type: "repeat", Detail: detailFromEntry(e)})

		case "intrude":
			if s.intruder == nil {
				enc.Encode(Response{Type: "error", Error: "intruder not available"})
				continue
			}
			err := s.intruder.Intrude(context.Background(), req.Scheme, req.Host, req.Raw, req.Payloads,
				func(position int, payload string, entry *store.Entry, sendErr error) bool {
					r := IntrudeResultMsg{Position: position, Payload: payload}
					if sendErr != nil {
						r.Error = sendErr.Error()
					}
					if entry != nil {
						r.EntryID = entry.ID
						r.StatusCode = entry.StatusCode
						r.RespSize = len(entry.ResponseRaw)
						r.Duration = entry.Duration
						if entry.Error != "" && r.Error == "" {
							r.Error = entry.Error
						}
					}
					return enc.Encode(Response{Type: "intrude_result", IntrudeResult: &r}) == nil
				})
			if err != nil {
				enc.Encode(Response{Type: "error", Error: err.Error()})
				continue
			}
			enc.Encode(Response{Type: "intrude_done"})

		case "rules_list":
			rs, err := s.db.ListRules()
			if err != nil {
				enc.Encode(Response{Type: "error", Error: err.Error()})
				continue
			}
			enc.Encode(Response{Type: "rules", Rules: rs})

		case "rules_save":
			if req.Rule == nil {
				enc.Encode(Response{Type: "error", Error: "rules_save: missing rule"})
				continue
			}
			r := *req.Rule
			var err error
			if r.ID == 0 {
				r.ID, err = s.db.AddRule(r)
			} else {
				err = s.db.UpdateRule(r)
			}
			if err != nil {
				enc.Encode(Response{Type: "error", Error: err.Error()})
				continue
			}
			enc.Encode(Response{Type: "rules", Rules: []rules.Rule{r}})

		case "rules_delete":
			if err := s.db.DeleteRule(req.RuleID); err != nil {
				enc.Encode(Response{Type: "error", Error: err.Error()})
				continue
			}
			enc.Encode(Response{Type: "rules"})

		case "rules_toggle":
			if err := s.db.SetRuleEnabled(req.RuleID, req.RuleEnabled); err != nil {
				enc.Encode(Response{Type: "error", Error: err.Error()})
				continue
			}
			enc.Encode(Response{Type: "rules"})

		case "subscribe":
			sub := s.hub.subscribe()
			defer s.hub.unsubscribe(sub)

			// The client never sends anything more on this connection,
			// but reading it anyway is how we notice it went away when
			// no new entry ever arrives to trigger a failed Encode -
			// otherwise a subscriber that quits while the daemon is
			// idle leaks its goroutine and channel indefinitely.
			disconnected := make(chan struct{})
			go func() {
				io.Copy(io.Discard, conn)
				close(disconnected)
			}()

			for {
				select {
				case e, ok := <-sub:
					if !ok {
						return
					}
					if err := enc.Encode(Response{Type: "new", New: &e}); err != nil {
						return
					}
				case <-disconnected:
					return
				}
			}

		default:
			enc.Encode(Response{Type: "error", Error: "unknown request type: " + req.Type})
		}
	}
}

func detailFromEntry(e *store.Entry) *EntryDetail {
	return &EntryDetail{
		Summary: store.Summary{
			ID: e.ID, StartedAt: e.StartedAt, Duration: e.Duration,
			Method: e.Method, Scheme: e.Scheme, Host: e.Host, Path: e.Path,
			StatusCode: e.StatusCode, ReqSize: len(e.RequestRaw), RespSize: len(e.ResponseRaw),
			Error: e.Error, Source: e.Source,
		},
		RequestRaw: e.RequestRaw, ResponseRaw: e.ResponseRaw,
		RequestExact: e.RequestExact, ResponseExact: e.ResponseExact,
	}
}

// LogAndBroadcast is a convenience OnEntry callback: logs the entry and
// broadcasts it through hub.
func LogAndBroadcast(hub *Hub) func(store.Summary) {
	return func(sum store.Summary) {
		log.Printf("%s %s%s -> %d (%s)", sum.Method, sum.Host, sum.Path, sum.StatusCode, sum.Duration)
		hub.Broadcast(sum)
	}
}