proton-bridge/test/mocks/imap_response.go

// Copyright (c) 2020 Proton Technologies AG
//
// This file is part of ProtonMail Bridge.Bridge.
//
// ProtonMail Bridge is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// ProtonMail Bridge is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with ProtonMail Bridge.  If not, see <https://www.gnu.org/licenses/>.

package mocks

import (
	"fmt"
	"io"
	"regexp"
	"strings"
	"time"

	"github.com/emersion/go-imap"
	"github.com/pkg/errors"
	a "github.com/stretchr/testify/assert"
)

type IMAPResponse struct {
	t        TestingT
	err      error
	result   string
	sections []string
	done     bool
}

func (ir *IMAPResponse) sendCommand(reqTag string, reqIndex int, command string, debug *debug, conn io.Writer, response imap.StringReader) {
	defer func() { ir.done = true }()

	tstart := time.Now()

	commandID := fmt.Sprintf("%sO%0d", reqTag, reqIndex)
	command = fmt.Sprintf("%s %s", commandID, command)

	debug.printReq(command)
	fmt.Fprintf(conn, "%s\r\n", command)

	var section string
	for {
		line, err := response.ReadString('\n')

		if err != nil {
			ir.err = errors.Wrap(err, "read response failed")
			debug.printErr(ir.err.Error() + "\n")
			return
		}

		// Finishing line contains `commandID` following with status (`NO`, `BAD`, ...) and then message itself.
		lineWithoutID := strings.Replace(line, commandID+" ", "", 1)
		if strings.HasPrefix(line, commandID) && (strings.HasPrefix(lineWithoutID, "NO ") || strings.HasPrefix(lineWithoutID, "BAD ")) {
			debug.printErr(line)
			err := errors.New(strings.Trim(lineWithoutID, "\r\n"))
			ir.err = errors.Wrap(err, "IMAP error")
			return
		} else if command != "" && len(line) == 0 {
			err := errors.New("empty answer")
			ir.err = errors.Wrap(err, "IMAP error")
			debug.printErr(ir.err.Error() + "\n")
			return
		}
		debug.printRes(line)

		if strings.HasPrefix(line, "* ") { //nolint[gocritic]
			if section != "" {
				ir.sections = append(ir.sections, section)
			}
			section = line
		} else if strings.HasPrefix(line, commandID) {
			if section != "" {
				ir.sections = append(ir.sections, section)
			}
			ir.result = line
			break
		} else {
			section += line
		}
	}

	debug.printTime(time.Since(tstart))
}

func (ir *IMAPResponse) wait() {
	for {
		if ir.done {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}
}

func (ir *IMAPResponse) AssertOK() *IMAPResponse {
	ir.wait()
	a.NoError(ir.t, ir.err)
	return ir
}

func (ir *IMAPResponse) AssertError(wantErrMsg string) *IMAPResponse {
	ir.wait()
	if ir.err == nil {
		a.Fail(ir.t, "Expected error %s", wantErrMsg)
	} else {
		a.Regexp(ir.t, wantErrMsg, ir.err.Error(), "Expected error %s but got %s", wantErrMsg, ir.err)
	}
	return ir
}

func (ir *IMAPResponse) AssertSectionsCount(expectedCount int) *IMAPResponse {
	ir.wait()
	a.Equal(ir.t, expectedCount, len(ir.sections))
	return ir
}

// AssertSectionsInOrder checks sections against regular expression in exact order.
// First regexp checks first section, second the second and so on. If there is
// more responses (sections) than expected regexps, that's OK.
func (ir *IMAPResponse) AssertSectionsInOrder(wantRegexps ...string) *IMAPResponse {
	ir.wait()
	if !a.True(ir.t,
		len(ir.sections) >= len(wantRegexps),
		"Wrong number of sections, want %v, got %v",
		len(wantRegexps),
		len(ir.sections),
	) {
		return ir
	}

	for idx, wantRegexp := range wantRegexps {
		section := ir.sections[idx]
		match, err := regexp.MatchString(wantRegexp, section)
		if !a.NoError(ir.t, err) {
			return ir
		}
		if !a.True(ir.t, match, "Section does not match given regex", section, wantRegexp) {
			return ir
		}
	}
	return ir
}

// AssertSections is similar to AssertSectionsInOrder but is not strict to the order.
// It means it just tries to find all "regexps" in the response.
func (ir *IMAPResponse) AssertSections(wantRegexps ...string) *IMAPResponse {
	ir.wait()
	for _, wantRegexp := range wantRegexps {
		a.NoError(ir.t, ir.hasSectionRegexp(wantRegexp), "regexp %v not found", wantRegexp)
	}
	return ir
}

// WaitForSections is the same as AssertSections but waits for `timeout` before giving up.
func (ir *IMAPResponse) WaitForSections(timeout time.Duration, wantRegexps ...string) {
	a.Eventually(ir.t, func() bool {
		return ir.HasSections(wantRegexps...)
	}, timeout, 50*time.Millisecond, "Wanted sections: %v\nSections: %v", wantRegexps, &ir.sections)
}

// WaitForNotSections is the opposite of WaitForSection: waits to not have the response.
func (ir *IMAPResponse) WaitForNotSections(timeout time.Duration, unwantedRegexps ...string) *IMAPResponse {
	time.Sleep(timeout)
	match := ir.HasSections(unwantedRegexps...)
	a.False(ir.t, match, "Unwanted sections: %v\nSections: %v", unwantedRegexps, &ir.sections)
	return ir
}

// HasSections is the same as AssertSections but only returns bool (do not uses testingT).
func (ir *IMAPResponse) HasSections(wantRegexps ...string) bool {
	for _, wantRegexp := range wantRegexps {
		if err := ir.hasSectionRegexp(wantRegexp); err != nil {
			return false
		}
	}
	return true
}

func (ir *IMAPResponse) hasSectionRegexp(wantRegexp string) error {
	for _, section := range ir.sections {
		match, err := regexp.MatchString(wantRegexp, section)
		if err != nil {
			return err
		}
		if match {
			return nil
		}
	}
	return errors.New("Section matching given regex not found")
}