grpc-go/internal/resolver/dns/fake_net_resolver_test.go

127 lines
3.5 KiB
Go

/*
*
* Copyright 2023 gRPC authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
package dns_test
import (
"context"
"net"
"sync"
"google.golang.org/grpc/internal/testutils"
)
// A fake implementation of the internal.NetResolver interface for use in tests.
type testNetResolver struct {
// A write to this channel is made when this resolver receives a resolution
// request. Tests can rely on reading from this channel to be notified about
// resolution requests instead of sleeping for a predefined period of time.
lookupHostCh *testutils.Channel
mu sync.Mutex
hostLookupTable map[string][]string // Name --> list of addresses
srvLookupTable map[string][]*net.SRV // Name --> list of SRV records
txtLookupTable map[string][]string // Name --> service config for TXT record
}
func (tr *testNetResolver) LookupHost(ctx context.Context, host string) ([]string, error) {
if tr.lookupHostCh != nil {
if err := tr.lookupHostCh.SendContext(ctx, nil); err != nil {
return nil, err
}
}
tr.mu.Lock()
defer tr.mu.Unlock()
if addrs, ok := tr.hostLookupTable[host]; ok {
return addrs, nil
}
return nil, &net.DNSError{
Err: "hostLookup error",
Name: host,
Server: "fake",
IsTemporary: true,
}
}
func (tr *testNetResolver) UpdateHostLookupTable(table map[string][]string) {
tr.mu.Lock()
tr.hostLookupTable = table
tr.mu.Unlock()
}
func (tr *testNetResolver) LookupSRV(_ context.Context, service, proto, name string) (string, []*net.SRV, error) {
tr.mu.Lock()
defer tr.mu.Unlock()
cname := "_" + service + "._" + proto + "." + name
if srvs, ok := tr.srvLookupTable[cname]; ok {
return cname, srvs, nil
}
return "", nil, &net.DNSError{
Err: "srvLookup error",
Name: cname,
Server: "fake",
IsTemporary: true,
}
}
func (tr *testNetResolver) LookupTXT(_ context.Context, host string) ([]string, error) {
tr.mu.Lock()
defer tr.mu.Unlock()
if sc, ok := tr.txtLookupTable[host]; ok {
return sc, nil
}
return nil, &net.DNSError{
Err: "txtLookup error",
Name: host,
Server: "fake",
IsTemporary: true,
}
}
func (tr *testNetResolver) UpdateTXTLookupTable(table map[string][]string) {
tr.mu.Lock()
tr.txtLookupTable = table
tr.mu.Unlock()
}
// txtRecordServiceConfig generates a slice of strings (aggregately representing
// a single service config file) for the input config string, that represents
// the result from a real DNS TXT record lookup.
func txtRecordServiceConfig(cfg string) []string {
// In DNS, service config is encoded in a TXT record via the mechanism
// described in RFC-1464 using the attribute name grpc_config.
b := append([]byte("grpc_config="), []byte(cfg)...)
// Split b into multiple strings, each with a max of 255 bytes, which is
// the DNS TXT record limit.
var r []string
for i := 0; i < len(b); i += txtBytesLimit {
if i+txtBytesLimit > len(b) {
r = append(r, string(b[i:]))
} else {
r = append(r, string(b[i:i+txtBytesLimit]))
}
}
return r
}