crawler.go

package main

import (
	"errors"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/html"
)

// -------- Class declaration and constructor -------- //

type Crawler struct {
	data_handler  DataHandler
	stop_words    map[string]struct{}
	permissions   map[string]bool
	delays        map[string]float32
	default_delay float32
}

// Constructor for an empty crawler with a set of stop words given by the file path
func NewCrawler(data_handler DataHandler, stop_words map[string]struct{}, default_delay float32) *Crawler {
	// Create crawler
	return &Crawler{
		data_handler:  data_handler,
		stop_words:    stop_words,
		permissions:   make(map[string]bool),
		delays:        make(map[string]float32),
		default_delay: default_delay,
	}
}

// -------- Core Crawler Components -------- //

// Cleans hrefs by converting them to their full url instead of relative forms
func (c *Crawler) CleanHref(cur_url string, href string) (string, error) {
	// If href is relative, join with root of parent url
	var abs_href string
	_, err := GetDomain(href)
	if err != nil {
		cur_domain, err := GetDomain(cur_url)
		if err != nil {
			return href, err
		}
		joined_href, err := url.JoinPath(cur_domain, href)
		if err != nil {
			return href, errors.New("unable to join href with root url")
		}
		abs_href = joined_href
	} else {
		abs_href = href
	}

	// Normalize absolute href
	normal_href, err := url.Parse(abs_href)
	if err != nil {
		return abs_href, err
	}
	return normal_href.String(), nil
}

// Downloads html from the given url
func (c *Crawler) Download(url string) ([]byte, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}

	html_data, err := io.ReadAll(resp.Body)
	defer resp.Body.Close()
	if err != nil {
		return nil, err
	}
	return html_data, nil
}

// Takes in raw html data and returns all the words (non-duplicating) and non-absolute hrefs within
func (c *Crawler) Extract(rawdata []byte) (string, map[string]uint, map[string]struct{}, uint) {
	// Define collections
	title := ""
	word_map := make(map[string]uint)
	href_subset := make(map[string]struct{})
	word_count := uint(0)

	// Create word collector
	getWords := func(text string) {
		words := NormalizeWords(word_re.FindAllString(text, -1))
		for _, word := range words {
			if _, ok := c.stop_words[word]; !ok {
				word_map[word]++
			}
			word_count++
		}
	}

	// Create non-text tag filter
	badtags := map[string]struct{}{
		"head":     {},
		"link":     {},
		"meta":     {},
		"style":    {},
		"script":   {},
		"noscript": {},
	}

	// Use tokenizer to traverse html
	tokenizer := html.NewTokenizer(strings.NewReader(string(rawdata)))
	prev_tag := ""
	for {
		tt := tokenizer.Next()
		switch tt {
		// non-destructive error
		case html.ErrorToken:
			goto END
		// update to know when text is displayed
		case html.StartTagToken:
			token := tokenizer.Token()
			prev_tag = token.Data
			// get hrefs from anchors
			if token.Data == "a" {
				for _, a := range token.Attr {
					if a.Key == "href" {
						href_subset[a.Val] = struct{}{}
						break
					}
				}
			}
		// reset prev_token
		case html.EndTagToken, html.SelfClosingTagToken:
			prev_tag = ""
		// collect words from text
		case html.TextToken:
			_, ok := badtags[prev_tag]
			if ok {
				continue
			}
			text := tokenizer.Token().Data
			if title == "" && prev_tag == "title" {
				title = text
			}
			getWords(text)
		}
	}
END:
	return title, word_map, href_subset, word_count
}

// -------- Concurrency Logic -------- //

// Spawns workers that consume from raw_href_ch and
// serve href_ch by cleaning the consumed raw hrefs
func (c *Crawler) cleanHrefRoutine(
	max_workers int,
	raw_href_ch *QueueChannel[struct {
		prev string
		href string
	}],
	href_ch *QueueChannel[string],
	seen_map map[string]struct{},
	seen_lock *sync.Mutex,
	wait_group *sync.WaitGroup,
) {
	// Create worker manager
	throttle := max_workers > 0
	var time_sheet chan struct{}
	if throttle {
		time_sheet = make(chan struct{}, max_workers)
	}

	for {
		// Wait for html queue
		raw_href_zip := raw_href_ch.Pop()

		// Clock worker in
		if throttle {
			time_sheet <- struct{}{}
		}

		// Extract logic
		go func(
			throttle bool,
			time_sheet chan struct{},
			wait_group *sync.WaitGroup,
			zip struct {
				prev string
				href string
			},
		) {
			// Defer clock out
			defer func() {
				if throttle {
					<-time_sheet
				}
			}()

			// Clean href
			clean_href, err := c.CleanHref(zip.prev, zip.href)
			if err != nil {
				wait_group.Done()
				return
			}

			// Check crawl rule set
			allowed := true
			for rule, perm := range c.permissions {
				if !regexp.MustCompile(rule).MatchString(clean_href) {
					continue
				}
				allowed = perm
				if perm {
					break
				}
			}
			if !allowed {
				wait_group.Done()
				return
			}

			// Check confinement
			sameRoot := func(a string, b string) bool {
				domain_a, err_a := GetDomain(a)
				domain_b, err_b := GetDomain(b)
				if err_a != nil || err_b != nil {
					return false
				}
				return domain_a == domain_b
			}
			if !sameRoot(zip.prev, clean_href) {
				wait_group.Done()
				return
			}

			// Check seen for duplicate documents
			seen_lock.Lock()
			if _, ok := seen_map[clean_href]; ok {
				seen_lock.Unlock()
				wait_group.Done()
				return
			}
			seen_map[clean_href] = struct{}{}
			seen_lock.Unlock()

			// Add clean href to href channel
			go href_ch.Push(clean_href)
		}(throttle, time_sheet, wait_group, raw_href_zip)
	}
}

// Spawns workers that consume from href_ch and serve
// html_ch by downloading html from the consumed hrefs
func (c *Crawler) downloadRoutine(
	max_workers int,
	href_ch *QueueChannel[string],
	html_ch *QueueChannel[struct {
		href     string
		raw_data []byte
	}],
	crawl_delay float32,
	err_counter *int,
	err_lock *sync.Mutex,
	wait_group *sync.WaitGroup,
) {
	// Create worker manager
	throttle := max_workers > 0
	var time_sheet chan struct{}
	if throttle {
		time_sheet = make(chan struct{}, max_workers)
	}

	// Create delay manager
	var tracker_lock sync.Mutex
	tracker_time := time.Now().Add(-time.Duration(float32(crawl_delay) * float32(time.Second)))

	for {
		// Wait for href queue
		href := href_ch.Pop()

		// Clock worker in
		if throttle {
			time_sheet <- struct{}{}
		}

		// Download logic
		go func(
			throttle bool,
			time_sheet chan struct{},
			wait_group *sync.WaitGroup,
			href string,
		) {
			// Defer clock out
			defer func() {
				if throttle {
					<-time_sheet
				}
			}()

			// Find elapsed time and delay accordingly
			for {
				// Get elapsed time since last download
				tracker_lock.Lock()
				t_start := tracker_time
				elapsed := float32(time.Since(t_start).Seconds())

				// Wait until delay > expected delay
				if elapsed < crawl_delay {
					tracker_lock.Unlock()
					wait := crawl_delay - elapsed
					time.Sleep(time.Duration(float64(wait) * float64(time.Second)))
				} else {
					tracker_time = time.Now()
					tracker_lock.Unlock()
					break
				}
			}

			// Attempt 3 times max to download data
			var raw_data []byte
			var err error
			for range 3 {
				raw_data, err = c.Download(href)
				if err == nil {
					break
				} else if _, ok := err.(*url.Error); ok {
					continue
				} else {
					return
				}
			}
			if err != nil {
				err_lock.Lock()
				*err_counter++
				err_lock.Unlock()
				wait_group.Done()
				return
			}

			// Add html data to html channel
			go html_ch.Push(struct {
				href     string
				raw_data []byte
			}{
				href:     href,
				raw_data: raw_data,
			})
		}(throttle, time_sheet, wait_group, href)
	}
}

// Spawns workers that consume from html_ch and serve word_ch
// and raw_href_channel by extracting consumed html data
func (c *Crawler) extractRoutine(
	max_workers int,
	html_ch *QueueChannel[struct {
		href     string
		raw_data []byte
	}],
	raw_href_ch *QueueChannel[struct {
		prev string
		href string
	}],
	word_ch *QueueChannel[struct {
		title      string
		href       string
		word_count uint
		words      map[string]uint
	}],
	wait_group *sync.WaitGroup,
) {
	// Create worker manager
	throttle := max_workers > 0
	var time_sheet chan struct{}
	if throttle {
		time_sheet = make(chan struct{}, max_workers)
	}

	for {
		// Wait for html queue
		html_zip := html_ch.Pop()

		// Clock worker in
		if throttle {
			time_sheet <- struct{}{}
		}

		// Extract logic
		go func(
			throttle bool,
			time_sheet chan struct{},
			wait_group *sync.WaitGroup,
			zip struct {
				href     string
				raw_data []byte
			},
		) {
			// Defer clock out
			defer func() {
				if throttle {
					<-time_sheet
				}
			}()

			// Extract data from the html data
			title, words, raw_hrefs, word_count := c.Extract(zip.raw_data)
			if title == "" {
				title = zip.href
			}

			// Add raw hrefs to raw href channel
			for raw_href := range raw_hrefs {
				wait_group.Add(1)
				raw_href_ch.Push(struct {
					prev string
					href string
				}{
					prev: zip.href,
					href: raw_href,
				})
			}

			// Add words to words channel
			go word_ch.Push(struct {
				title      string
				href       string
				word_count uint
				words      map[string]uint
			}{
				title:      title,
				href:       zip.href,
				word_count: word_count,
				words:      words,
			})
		}(throttle, time_sheet, wait_group, html_zip)
	}
}

func (c *Crawler) indexRoutine(
	word_ch *QueueChannel[struct {
		title      string
		href       string
		word_count uint
		words      map[string]uint
	}],
	wait_group *sync.WaitGroup,
) {
	for {
		// Wait for word queue
		words_zip := word_ch.Pop()

		// Attempt 3 times to index document with words
		for range 3 {
			err := c.data_handler.AddDocWithWords(
				words_zip.title,
				words_zip.href,
				words_zip.word_count,
				words_zip.words,
			)
			if err == nil {
				break
			}
		}
		wait_group.Done()
	}
}

// -------- Crawl Logic -------- //

// Crawls a site and returns the reachable links with all words within them
func (c *Crawler) Crawl(url string, workers struct {
	clean_href int
	download   int
	extract    int
}) (int, error) {
	// Early stop in case of redundant crawl
	ok, err := c.data_handler.ContainsDoc(url)
	if err == nil {
		if ok {
			return 0, nil
		}
	} else {
		return 0, err
	}

	// Create crawl ruleset for this domain and get delay and sitemaps
	var crawl_delay float32
	var sitemaps []string
	domain, err := GetDomain(url)
	if err == nil {
		sitemaps = c.FindRulesAndSitemaps(domain)
		crawl_delay = c.delays[domain]
	} else {
		sitemaps = []string{}
		crawl_delay = c.default_delay
	}

	// Set up seen map by filling with known docs
	var seen_lock sync.Mutex
	seen_map := make(map[string]struct{})
	docs_map, err := c.data_handler.GetAllDocs()
	if err != nil {
		return 0, err
	}
	for doc := range docs_map {
		seen_map[doc] = struct{}{}
	}

	// Set up seen map and seen map lock
	var err_lock sync.Mutex
	err_counter := 0

	// Set up wait group and queue channels
	wait_group := sync.WaitGroup{}
	href_ch := NewQueueChannel[string]()
	html_ch := NewQueueChannel[struct {
		href     string
		raw_data []byte
	}]()
	raw_href_ch := NewQueueChannel[struct {
		prev string
		href string
	}]()
	word_ch := NewQueueChannel[struct {
		title      string
		href       string
		word_count uint
		words      map[string]uint
	}]()

	// Set up go routines
	go c.cleanHrefRoutine(
		workers.clean_href,
		raw_href_ch,
		href_ch,
		seen_map,
		&seen_lock,
		&wait_group,
	)
	go c.downloadRoutine(
		workers.download,
		href_ch,
		html_ch,
		crawl_delay,
		&err_counter,
		&err_lock,
		&wait_group,
	)
	go c.extractRoutine(
		workers.extract,
		html_ch,
		raw_href_ch,
		word_ch,
		&wait_group,
	)
	go c.indexRoutine(
		word_ch,
		&wait_group,
	)

	// Seed the href channel with initial url and sitemaps
	seen_map[url] = struct{}{}
	wait_group.Add(1)
	href_ch.Push(url)
	for _, site := range sitemaps {
		seen_map[site] = struct{}{}
		wait_group.Add(1)
		href_ch.Push(site)
	}

	// Wait for go routines to finish
	wait_group.Wait()
	return err_counter, nil
}

// Creates the crawling ruleset given a robots.txt file
func (c *Crawler) FindRulesAndSitemaps(domain string) []string {
	// Find and download robots.txt
	raw_text, err := c.Download(domain + "/robots.txt")
	if err != nil {
		// robots.txt does not exist
		c.permissions[domain+"/"] = true
		c.delays[domain] = c.default_delay
		return []string{}
	}

	// Find sitemaps
	sitemaps := []string{}
	sitemap_re := regexp.MustCompile(`(?m)^Sitemap: .+$`)
	sitemap_texts := sitemap_re.FindAllString(string(raw_text), -1)
	for _, sitemap_text := range sitemap_texts {
		split := strings.Split(strings.TrimSpace(sitemap_text), ": ")
		if len(split) == 2 {
			sitemaps = append(sitemaps, split[1])
		}
	}

	// Find only section for all user-agents
	all_agent_re := regexp.MustCompile(`(?s)User-agent: \*(.*?)(?:User-agent:|$)`)
	all_agent_texts := all_agent_re.FindStringSubmatch(string(raw_text))
	if len(all_agent_texts) < 2 {
		// robots.txt does not define limits for all agents
		c.permissions[domain+"/"] = true
		c.delays[domain] = c.default_delay
		return sitemaps
	}

	// Interpret robots.txt text for all agents
	rules := strings.Split(strings.TrimSpace(all_agent_texts[1]), "\n")
	for _, rule := range rules {
		// Parse rule into rule type and parameter
		parsed_rule := strings.Split(strings.TrimSpace(rule), ": ")
		if len(parsed_rule) != 2 {
			// Ignore rule if the rule was formatted incorrectly
			continue
		}
		command := strings.TrimSpace(parsed_rule[0])
		param := strings.TrimSpace(parsed_rule[1])

		// Add specified crawl delay for domain
		if command == "Crawl-delay" {
			if num, err := strconv.ParseFloat(param, 32); err == nil {
				c.delays[domain] = float32(num)
			} else {
				// Fallback to a default delay if parsing fails
				c.delays[domain] = c.default_delay
			}
			continue
		}

		// Add domain to front
		if !regexp.MustCompile(`^/`).MatchString(param) {
			param = "/" + param
		}
		param = domain + param

		// Normalize the constructed rule URL so it matches how CleanHref/url.Parse normalizes
		// hrefs used elsewhere (e.g., Remove default ports, resolve dots, trailing slashes)
		if parsed, err := url.Parse(param); err == nil {
			param = parsed.String()
		}

		// Switch url match to regex syntax
		param = regexp.MustCompile(`\.`).ReplaceAllString(param, `\.`)
		param = regexp.MustCompile(`\?`).ReplaceAllString(param, `\?`)
		param = regexp.MustCompile(`\*`).ReplaceAllString(param, `.*`)

		// Handle allow and disallow commands
		switch command {
		case "Allow":
			c.permissions[param] = true
		case "Disallow":
			c.permissions[param] = false
		}
	}
	return sitemaps
}