path: root/main.go

// Package main provides onion-grab, a tool that visits a list of domains over
// HTTPS to see if they have Onion-Location configured.
package main

import (
	"bufio"
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"git.cs.kau.se/rasmoste/onion-grab/internal/line"
	"git.cs.kau.se/rasmoste/onion-grab/internal/onionloc"
	"git.cs.kau.se/rasmoste/onion-grab/internal/options"
	"git.cs.kau.se/rasmoste/onion-grab/internal/qna"
)

func main() {
	start := time.Now().Round(time.Second)
	defer func() {
		end := time.Now().Round(time.Second)
		log.Printf("INFO: measurement duration was %v\n", end.Sub(start))
	}()

	opts := options.Parse()
	cli := &http.Client{
		Transport: &http.Transport{
			DisableKeepAlives:      true,
			MaxResponseHeaderBytes: opts.MaxResponse * 1024 * 1024,
		},
	}
	fp, err := os.OpenFile(opts.InputFile, os.O_RDONLY, 0644)
	if err != nil {
		log.Printf("ERROR: %v", err)
		os.Exit(1)
	}
	defer fp.Close()

	questionCh := make(chan qna.Question)
	defer close(questionCh)

	answerCh := make(chan qna.Answer)
	defer close(answerCh)

	var wg sync.WaitGroup
	defer wg.Wait()

	bg := context.Background()
	ctx, cancel := context.WithCancel(bg)
	defer cancel()

	log.Printf("INFO: ctrl+C to exit prematurely\n")
	go func() {
		wg.Add(1)
		defer wg.Done()
		await(ctx, cancel)
	}()

	log.Printf("INFO: starting %d workers\n", opts.Limit)
	for i := 0; i < opts.NumWorkers; i++ {
		go func() {
			wg.Add(1)
			defer wg.Done()
			for {
				select {
				case <-ctx.Done():
					return
				case question := <-questionCh:
					work(bg, cli, opts.Timeout, question, answerCh)
				}
			}
		}()
	}

	log.Printf("INFO: starting work receiver\n")
	go func() {
		wg.Add(1)
		defer wg.Done()
		workReceiver(ctx, opts, answerCh)
	}()

	log.Printf("INFO: starting work generator\n")
	workGenerator(ctx, opts, fp, questionCh)
}

func await(ctx context.Context, cancel context.CancelFunc) {
	sigs := make(chan os.Signal, 1)
	defer close(sigs)

	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	select {
	case <-sigs:
		log.Printf("WARNING: premature shutdown (context cancelled)\n")
		cancel()
	case <-ctx.Done():
	}
}

func work(ctx context.Context, cli *http.Client, timeout time.Duration, question qna.Question, answerCh chan qna.Answer) {
	cctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	answer := qna.Answer{Domain: question.Domain}
	req, err := http.NewRequestWithContext(cctx, http.MethodGet, "https://"+question.Domain, nil)
	if err != nil {
		answer.ReqErr = err
		answerCh <- answer
		return
	}

	rsp, err := cli.Do(req)
	if err != nil {
		answer.DoErr = err
		answerCh <- answer
		return
	}
	defer rsp.Body.Close()

	onion, ok := onionloc.HTTP(rsp)
	if ok {
		answer.HTTP = onion
	}
	onion, ok = onionloc.HTML(rsp)
	if ok {
		answer.HTML = onion
	}
	answerCh <- answer
}

func workReceiver(ctx context.Context, opts options.Options, answerCh chan qna.Answer) {
	p := qna.Progress{}
	handleAnswer := func(a qna.Answer) {
		p.AddAnswer(a)
		if a.HTTP != "" || a.HTML != "" {
			fmt.Printf("%s\n", a.String())
		}
	}

	metrics := time.NewTicker(opts.MetricsInterval)
	defer metrics.Stop()
	for {
		select {
		case <-ctx.Done():
			log.Printf("INFO: about to exit, reading remaining answers\n")
			for {
				select {
				case a := <-answerCh:
					handleAnswer(a)
				case <-time.After(opts.Timeout + time.Second):
					log.Printf("INFO: metrics@receiver: summary: \n\n%s\n\n", p.String())
					return
				}
			}
		case a := <-answerCh:
			handleAnswer(a)
		case <-metrics.C:
			log.Printf("INFO: metrics@receiver: \n\n%s\n\n", p.String())
		}
	}
}

func workGenerator(ctx context.Context, opts options.Options, fp *os.File, questionCh chan qna.Question) {
	var wg sync.WaitGroup
	defer wg.Wait()

	cctx, cancel := context.WithCancel(ctx)
	defer cancel()

	var nextLine line.Line
	go func() {
		wg.Add(1)
		defer wg.Done()
		generatorMetrics(cctx, opts, &nextLine)

		//
		// Would be nice to clean this up so that the Line type with a
		// mutex can be eliminated; and with all metrics in one place.
		//
	}()

	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	numRequests := 0

	scanner := bufio.NewScanner(fp)
	buf := make([]byte, 0, 128*1024*1024)
	scanner.Buffer(buf, 128*1024*1024)
	for scanner.Scan() {
		if numRequests == opts.Limit {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				numRequests = 0
			}
		}

		select {
		case <-ctx.Done():
			return
		case questionCh <- qna.Question{Domain: qna.TrimWildcard(scanner.Text())}:
			nextLine.Inc()
			numRequests++
		}
	}

	select {
	case <-ctx.Done():
	case <-time.After(opts.Timeout + time.Second):
	}
}

func generatorMetrics(ctx context.Context, opts options.Options, nextLine *line.Line) {
	metrics := time.NewTicker(opts.MetricsInterval)
	defer metrics.Stop()

	startTime := time.Now().Unix()
	prevTime := startTime
	prevLine := int64(0)
	for {
		select {
		case <-ctx.Done():
			return
		case <-metrics.C:
			currLine := nextLine.Num()
			now := time.Now().Unix()
			str := fmt.Sprintf("  Current rate: %.1f sites/s\n", float64(currLine-prevLine)/float64(now-prevTime))
			str += fmt.Sprintf("  Average rate: %.1f sites/s\n", float64(currLine)/float64(now-startTime))
			str += fmt.Sprintf("     Next line: %d", currLine)
			log.Printf("INFO: metrics@generator:\n\n%s\n\n", str)

			prevTime = now
			prevLine = currLine
		}
	}
}