plcdns / plcdns.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/miekg/dns"
)

// DID Document structure
type DIDDocument struct {
	Context            []interface{}        `json:"@context"`
	ID                 string               `json:"id"`
	AlsoKnownAs        []string             `json:"alsoKnownAs"`
	VerificationMethod []VerificationMethod `json:"verificationMethod"`
	Service            []Service            `json:"service"`
}

type VerificationMethod struct {
	ID                 string `json:"id"`
	Type               string `json:"type"`
	Controller         string `json:"controller"`
	PublicKeyMultibase string `json:"publicKeyMultibase"`
}

type Service struct {
	ID              string `json:"id"`
	Type            string `json:"type"`
	ServiceEndpoint string `json:"serviceEndpoint"`
}

type PLCHandler struct {
	plcDirectory string
	cache        map[string]*CachedDID
}

type CachedDID struct {
	Document  *DIDDocument
	Timestamp time.Time
}

type QueryType int

const (
	QueryHandle QueryType = iota
	QueryPDS
	QueryPubKey
	QueryLabeler
	QueryInvalid
)

func NewPLCHandler(plcDirectory string) *PLCHandler {
	return &PLCHandler{
		plcDirectory: plcDirectory,
		cache:        make(map[string]*CachedDID),
	}
}

func (h *PLCHandler) fetchDIDDocument(did string) (*DIDDocument, error) {
	// Check cache first (5 minute TTL)
	if cached, exists := h.cache[did]; exists {
		if time.Since(cached.Timestamp) < 5*time.Minute {
			log.Printf("Cache hit for %s", did)
			return cached.Document, nil
		}
	}

	url := fmt.Sprintf("%s/%s", h.plcDirectory, did)
	log.Printf("Fetching DID document from: %s", url)

	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Get(url)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch DID document: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("DID not found: status %d", resp.StatusCode)
	}

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}

	var doc DIDDocument
	if err := json.Unmarshal(body, &doc); err != nil {
		return nil, fmt.Errorf("failed to parse DID document: %w", err)
	}

	// Cache the result
	h.cache[did] = &CachedDID{
		Document:  &doc,
		Timestamp: time.Now(),
	}

	return &doc, nil
}

// Parse domain name to extract DID and query type
// Expected formats:
// _handle.<did>.plc.atscan.net
// _pds.<did>.plc.atscan.net
// _pubkey.<did>.plc.atscan.net
// _labeler.<did>.plc.atscan.net
func (h *PLCHandler) parseDomain(domain string) (string, QueryType, bool) {
	domain = strings.TrimSuffix(domain, ".")
	parts := strings.Split(domain, ".")

	// Should be at least: [_prefix, <did>, plc, atscan, net]
	if len(parts) < 5 {
		return "", QueryInvalid, false
	}

	// Determine query type based on prefix
	var queryType QueryType
	switch parts[0] {
	case "_handle":
		queryType = QueryHandle
	case "_pds":
		queryType = QueryPDS
	case "_pubkey":
		queryType = QueryPubKey
	case "_labeler":
		queryType = QueryLabeler
	default:
		return "", QueryInvalid, false
	}

	// Extract DID identifier
	didIdentifier := parts[1]

	// Construct full DID
	did := fmt.Sprintf("did:plc:%s", didIdentifier)

	return did, queryType, true
}

// Get handle from DID document
func (h *PLCHandler) getHandle(doc *DIDDocument) string {
	for _, aka := range doc.AlsoKnownAs {
		if strings.HasPrefix(aka, "at://") {
			return strings.TrimPrefix(aka, "at://")
		}
	}
	return ""
}

// Get PDS endpoint from DID document
func (h *PLCHandler) getPDS(doc *DIDDocument) string {
	for _, service := range doc.Service {
		if service.Type == "AtprotoPersonalDataServer" {
			return service.ServiceEndpoint
		}
	}
	return ""
}

// Get labeler endpoint from DID document
func (h *PLCHandler) getLabeler(doc *DIDDocument) string {
	for _, service := range doc.Service {
		if service.ID == "#atproto_labeler" {
			return service.ServiceEndpoint
		}
	}
	return ""
}

// Get public key from DID document
func (h *PLCHandler) getPubKey(doc *DIDDocument) string {
	if len(doc.VerificationMethod) > 0 {
		return doc.VerificationMethod[0].PublicKeyMultibase
	}
	return ""
}

// Create TXT record based on query type
func (h *PLCHandler) createTXTRecord(doc *DIDDocument, qname string, queryType QueryType) []dns.RR {
	var records []dns.RR
	ttl := uint32(300) // 5 minutes

	var value string
	switch queryType {
	case QueryHandle:
		value = h.getHandle(doc)
		if value == "" {
			return records
		}
	case QueryPDS:
		value = h.getPDS(doc)
		if value == "" {
			return records
		}
	case QueryLabeler:
		value = h.getLabeler(doc)
		if value == "" {
			return records
		}
	case QueryPubKey:
		value = h.getPubKey(doc)
		if value == "" {
			return records
		}
	default:
		return records
	}

	records = append(records, &dns.TXT{
		Hdr: dns.RR_Header{
			Name:   qname,
			Rrtype: dns.TypeTXT,
			Class:  dns.ClassINET,
			Ttl:    ttl,
		},
		Txt: []string{value},
	})

	return records
}

func (h *PLCHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	msg := dns.Msg{}
	msg.SetReply(r)
	msg.Authoritative = true

	for _, question := range r.Question {
		log.Printf("Query: %s %s", question.Name, dns.TypeToString[question.Qtype])

		// Only handle TXT queries
		if question.Qtype != dns.TypeTXT {
			continue
		}

		// Parse the domain to extract DID and query type
		did, queryType, ok := h.parseDomain(question.Name)
		if !ok {
			log.Printf("Invalid domain format: %s", question.Name)
			msg.SetRcode(r, dns.RcodeNameError)
			w.WriteMsg(&msg)
			return
		}

		log.Printf("Extracted DID: %s, Query Type: %v", did, queryType)

		// Fetch DID document
		doc, err := h.fetchDIDDocument(did)
		if err != nil {
			log.Printf("Error fetching DID document: %s", err)
			msg.SetRcode(r, dns.RcodeServerFailure)
			w.WriteMsg(&msg)
			return
		}

		// Create TXT record based on query type
		records := h.createTXTRecord(doc, question.Name, queryType)
		if len(records) == 0 {
			log.Printf("No data found for query type")
			msg.SetRcode(r, dns.RcodeNameError)
			w.WriteMsg(&msg)
			return
		}

		msg.Answer = append(msg.Answer, records...)
	}

	if len(msg.Answer) == 0 {
		msg.SetRcode(r, dns.RcodeNameError)
	}

	w.WriteMsg(&msg)
}

func main() {
	// Command line flags
	port := flag.String("port", "", "DNS server port (default: 8053 or DNS_PORT env var)")
	plcDir := flag.String("plc", "https://plc.directory", "PLC directory URL")
	flag.Parse()

	// Determine port: flag -> env var -> default
	finalPort := *port
	if finalPort == "" {
		if envPort := os.Getenv("DNS_PORT"); envPort != "" {
			finalPort = envPort
		} else {
			finalPort = "8053"
		}
	}

	// Validate port number
	if portNum, err := strconv.Atoi(finalPort); err != nil || portNum < 1 || portNum > 65535 {
		log.Fatalf("Invalid port number: %s", finalPort)
	}

	addr := ":" + finalPort
	handler := NewPLCHandler(*plcDir)

	log.Printf("PLC Directory: %s", *plcDir)
	log.Printf("Starting DNS servers on port %s", finalPort)
	log.Printf("Supported query types:")
	log.Printf("  _handle.<did>.plc.atscan.net   - Returns handle")
	log.Printf("  _pds.<did>.plc.atscan.net      - Returns PDS endpoint")
	log.Printf("  _labeler.<did>.plc.atscan.net  - Returns labeler endpoint")
	log.Printf("  _pubkey.<did>.plc.atscan.net   - Returns public key")

	// UDP server
	udpServer := &dns.Server{
		Addr:    addr,
		Net:     "udp",
		Handler: handler,
	}

	// TCP server
	tcpServer := &dns.Server{
		Addr:    addr,
		Net:     "tcp",
		Handler: handler,
	}

	// Start UDP server in goroutine
	go func() {
		log.Printf("UDP DNS server listening on %s", addr)
		if err := udpServer.ListenAndServe(); err != nil {
			log.Fatalf("Failed to start UDP server: %s", err)
		}
	}()

	// Start TCP server
	log.Printf("TCP DNS server listening on %s", addr)
	if err := tcpServer.ListenAndServe(); err != nil {
		log.Fatalf("Failed to start TCP server: %s", err)
	}
}