google.golang.org/grpc@v1.62.1/internal/resolver/dns/fake_net_resolver_test.go (about)

     1  /*
     2   *
     3   * Copyright 2023 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  package dns_test
    20  
    21  import (
    22  	"context"
    23  	"net"
    24  	"sync"
    25  
    26  	"google.golang.org/grpc/internal/testutils"
    27  )
    28  
    29  // A fake implementation of the internal.NetResolver interface for use in tests.
    30  type testNetResolver struct {
    31  	// A write to this channel is made when this resolver receives a resolution
    32  	// request. Tests can rely on reading from this channel to be notified about
    33  	// resolution requests instead of sleeping for a predefined period of time.
    34  	lookupHostCh *testutils.Channel
    35  
    36  	mu              sync.Mutex
    37  	hostLookupTable map[string][]string   // Name --> list of addresses
    38  	srvLookupTable  map[string][]*net.SRV // Name --> list of SRV records
    39  	txtLookupTable  map[string][]string   // Name --> service config for TXT record
    40  }
    41  
    42  func (tr *testNetResolver) LookupHost(ctx context.Context, host string) ([]string, error) {
    43  	if tr.lookupHostCh != nil {
    44  		tr.lookupHostCh.Send(nil)
    45  	}
    46  
    47  	tr.mu.Lock()
    48  	defer tr.mu.Unlock()
    49  
    50  	if addrs, ok := tr.hostLookupTable[host]; ok {
    51  		return addrs, nil
    52  	}
    53  	return nil, &net.DNSError{
    54  		Err:         "hostLookup error",
    55  		Name:        host,
    56  		Server:      "fake",
    57  		IsTemporary: true,
    58  	}
    59  }
    60  
    61  func (tr *testNetResolver) UpdateHostLookupTable(table map[string][]string) {
    62  	tr.mu.Lock()
    63  	tr.hostLookupTable = table
    64  	tr.mu.Unlock()
    65  }
    66  
    67  func (tr *testNetResolver) LookupSRV(ctx context.Context, service, proto, name string) (string, []*net.SRV, error) {
    68  	tr.mu.Lock()
    69  	defer tr.mu.Unlock()
    70  
    71  	cname := "_" + service + "._" + proto + "." + name
    72  	if srvs, ok := tr.srvLookupTable[cname]; ok {
    73  		return cname, srvs, nil
    74  	}
    75  	return "", nil, &net.DNSError{
    76  		Err:         "srvLookup error",
    77  		Name:        cname,
    78  		Server:      "fake",
    79  		IsTemporary: true,
    80  	}
    81  }
    82  
    83  func (tr *testNetResolver) LookupTXT(ctx context.Context, host string) ([]string, error) {
    84  	tr.mu.Lock()
    85  	defer tr.mu.Unlock()
    86  
    87  	if sc, ok := tr.txtLookupTable[host]; ok {
    88  		return sc, nil
    89  	}
    90  	return nil, &net.DNSError{
    91  		Err:         "txtLookup error",
    92  		Name:        host,
    93  		Server:      "fake",
    94  		IsTemporary: true,
    95  	}
    96  }
    97  
    98  func (tr *testNetResolver) UpdateTXTLookupTable(table map[string][]string) {
    99  	tr.mu.Lock()
   100  	tr.txtLookupTable = table
   101  	tr.mu.Unlock()
   102  }
   103  
   104  // txtRecordServiceConfig generates a slice of strings (aggregately representing
   105  // a single service config file) for the input config string, that represents
   106  // the result from a real DNS TXT record lookup.
   107  func txtRecordServiceConfig(cfg string) []string {
   108  	// In DNS, service config is encoded in a TXT record via the mechanism
   109  	// described in RFC-1464 using the attribute name grpc_config.
   110  	b := append([]byte("grpc_config="), []byte(cfg)...)
   111  
   112  	// Split b into multiple strings, each with a max of 255 bytes, which is
   113  	// the DNS TXT record limit.
   114  	var r []string
   115  	for i := 0; i < len(b); i += txtBytesLimit {
   116  		if i+txtBytesLimit > len(b) {
   117  			r = append(r, string(b[i:]))
   118  		} else {
   119  			r = append(r, string(b[i:i+txtBytesLimit]))
   120  		}
   121  	}
   122  	return r
   123  }