-
Notifications
You must be signed in to change notification settings - Fork 177
/
options.go
130 lines (113 loc) · 3.95 KB
/
options.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
/*
* Copyright (c) 2019-2022. Abstrium SAS <team (at) pydio.com>
* This file is part of Pydio Cells.
*
* Pydio Cells is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* Pydio Cells 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 Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with Pydio Cells. If not, see <http://www.gnu.org/licenses/>.
*
* The latest code can be found at <https://pydio.com>.
*/
package client
import (
"net"
"os"
"strconv"
"strings"
"sync"
"google.golang.org/grpc/attributes"
"github.com/pydio/cells/v4/common/runtime"
)
type attKeyTargetServerID struct{}
type attKeyTargetServerPID struct{}
var (
localAddrOnce sync.Once
localAddress string
)
// BalancerTarget is a load balancer target providing its own internal info
type BalancerTarget interface {
// Address must return the target address. It can include a port
Address() string
// Attributes returns a set of grpc.Attributes
Attributes() *attributes.Attributes
}
// BalancerTargetFilter is a generic signature for matching a BalancerTarget against specific condition.
type BalancerTargetFilter func(BalancerTarget) bool
// BalancerOptions holds internal load balancer strategies
type BalancerOptions struct {
// Filters is a list of filters that will reduce the possible list of targets
Filters []BalancerTargetFilter
// Priority is a list of filters that will be tested first. If no target match, the load balancer
// continues to the non-matching values
Priority []BalancerTargetFilter
}
type BalancerOption func(o *BalancerOptions)
// WithRestrictToLocal provides a BalancerOption to forbid connection to hosts other than local.
func WithRestrictToLocal() BalancerOption {
return func(o *BalancerOptions) {
o.Filters = append(o.Filters, func(target BalancerTarget) bool {
return TargetHostMatches(target, getLocalAddress())
})
}
}
// WithPriorityToLocal provides a BalancerOption to prioritize connections to local host.
func WithPriorityToLocal() BalancerOption {
return func(o *BalancerOptions) {
o.Priority = append(o.Priority, func(target BalancerTarget) bool {
return TargetHostMatches(target, getLocalAddress())
})
}
}
// WithRestrictToProcess is mainly for testing purpose, forbidding connection to other processes
func WithRestrictToProcess() BalancerOption {
return func(o *BalancerOptions) {
o.Filters = append(o.Filters, TargetPIDMatches)
}
}
// WithPriorityToProcess is mainly for testing purpopse, to prioritize connections to current process.
func WithPriorityToProcess() BalancerOption {
return func(o *BalancerOptions) {
o.Priority = append(o.Priority, TargetPIDMatches)
}
}
// TargetHostMatches checks target.Address host against the passed string.
// hosts string is a pipe-separated list of hostname or IPs.
func TargetHostMatches(target BalancerTarget, hosts string) bool {
//fmt.Println("Check if we restrict to host here!", hosts, "vs", target.Address())
hh := strings.Split(hosts, "|")
aHost, _, er := net.SplitHostPort(target.Address())
if er != nil {
return false
}
for _, h := range hh {
if aHost == h {
return true
}
}
return false
}
// TargetPIDMatches is a generic matcher for PID
func TargetPIDMatches(target BalancerTarget) bool {
srvPID := target.Attributes().Value(attKeyTargetServerPID{})
//fmt.Println("Check if we restrict to PID here!", srvPID, "vs", os.Getpid())
if srvPID != nil {
localPID := strconv.Itoa(os.Getpid())
return srvPID.(string) == localPID
}
return false
}
func getLocalAddress() string {
localAddrOnce.Do(func() {
localAddress = runtime.DefaultAdvertiseAddress()
})
return localAddress
}