Passed
Pull Request — master (#52)
by Stefano
02:59
created

producer.*ReProducer.reprocess   B

Complexity

Conditions 7

Size

Total Lines 28
Code Lines 18

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 7
eloc 18
nop 1
dl 0
loc 28
rs 8
c 0
b 0
f 0
1
package producer
2
3
import (
4
	"path"
5
	"sync"
6
	"time"
7
8
	"github.com/stefanoj3/dirstalk/pkg/pathutil"
9
	"github.com/stefanoj3/dirstalk/pkg/scan"
10
)
11
12
const defaultChannelBuffer = 50
13
14
var statusCodesToSkip = map[int]bool{
15
	404: false,
16
}
17
18
func NewReProducer(
0 ignored issues
show
introduced by
exported function NewReProducer should have comment or be unexported
Loading history...
19
	httpMethods []string,
20
	dictionary []string,
21
	decoratedProducer scan.Producer,
22
	reproducerTimeout time.Duration,
23
) *ReProducer {
24
	return &ReProducer{
25
		httpMethods:       httpMethods,
26
		dictionary:        dictionary,
27
		decoratedProducer: decoratedProducer,
28
		reproducerTimeout: reproducerTimeout,
29
		reproducerChannel: make(chan scan.Target, defaultChannelBuffer),
30
	}
31
}
32
33
// ReProducer decorates a producer and adds to the produced Targets new targets based on
34
// result that must be provided via the Reproduce method.
35
// The ReProducer closes its "Produce" channel as soon as the decorated dictionary
36
// closes his channel and the last call to the Reproduce exceeds the timeout
37
type ReProducer struct {
38
	httpMethods       []string
39
	dictionary        []string
40
	decoratedProducer scan.Producer
41
	reproducerTimeout time.Duration
42
43
	reproducerChannel chan scan.Target
44
	resultRegistry    sync.Map
45
46
	lastTouchRW sync.RWMutex
47
	lastTouch   time.Time
48
}
49
50
func (r *ReProducer) Produce() <-chan scan.Target {
0 ignored issues
show
introduced by
exported method ReProducer.Produce should have comment or be unexported
Loading history...
51
	resultChannel := make(chan scan.Target, defaultChannelBuffer)
52
53
	decoratedProducerChannel := r.decoratedProducer.Produce()
54
55
	wg := sync.WaitGroup{}
56
57
	wg.Add(1)
58
	go func() {
59
		defer wg.Done()
60
		for producedTarget := range decoratedProducerChannel {
61
			resultChannel <- producedTarget
62
		}
63
	}()
64
65
	wg.Add(1)
66
	go func() {
67
		defer wg.Done()
68
		for {
69
			select {
70
			case reproducedTarget := <-r.reproducerChannel:
71
				resultChannel <- reproducedTarget
72
			case <-time.After(time.Millisecond * 100):
73
				r.lastTouchRW.RLock()
74
				diff := time.Since(r.lastTouch)
75
				r.lastTouchRW.RUnlock()
76
77
				if len(decoratedProducerChannel) == 0 && diff > r.reproducerTimeout {
78
					return
79
				}
80
			}
81
		}
82
	}()
83
84
	go func() {
85
		wg.Wait()
86
		close(resultChannel)
87
		close(r.reproducerChannel)
88
	}()
89
90
	return resultChannel
91
}
92
93
// Reproduce will check if it is possible to go deeper on the result provided, if so will queue a new target
94
func (r *ReProducer) Reproduce(result scan.Result) {
95
	go r.reprocess(result)
96
97
	r.lastTouchRW.Lock()
98
	defer r.lastTouchRW.Unlock()
99
100
	r.lastTouch = time.Now()
101
}
102
103
func (r *ReProducer) reprocess(result scan.Result) {
104
	if _, ok := statusCodesToSkip[result.Response.StatusCode]; ok {
105
		return
106
	}
107
108
	if result.Target.Depth <= 0 {
109
		return
110
	}
111
112
	// no point in appending to a filename
113
	if pathutil.HasExtension(result.Target.Path) {
114
		return
115
	}
116
117
	_, inRegistry := r.resultRegistry.Load(result.Target.Path)
118
	if inRegistry {
119
		return
120
	}
121
	r.resultRegistry.Store(result.Target.Path, false)
122
123
	for _, entry := range r.dictionary {
124
		for _, httpMethod := range r.httpMethods {
125
			t := result.Target
126
			t.Depth--
127
			t.Path = path.Join(t.Path, entry)
128
			t.Method = httpMethod
129
130
			r.reproducerChannel <- t
131
		}
132
	}
133
}
134