fake_net_resolver_test.go

Documentation: google.golang.org/grpc/internal/resolver/dns

     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  		if err := tr.lookupHostCh.SendContext(ctx, nil); err != nil {
    45  			return nil, err
    46  		}
    47  	}
    48  
    49  	tr.mu.Lock()
    50  	defer tr.mu.Unlock()
    51  
    52  	if addrs, ok := tr.hostLookupTable[host]; ok {
    53  		return addrs, nil
    54  	}
    55  
    56  	return nil, &net.DNSError{
    57  		Err:         "hostLookup error",
    58  		Name:        host,
    59  		Server:      "fake",
    60  		IsTemporary: true,
    61  	}
    62  }
    63  
    64  func (tr *testNetResolver) UpdateHostLookupTable(table map[string][]string) {
    65  	tr.mu.Lock()
    66  	tr.hostLookupTable = table
    67  	tr.mu.Unlock()
    68  }
    69  
    70  func (tr *testNetResolver) LookupSRV(ctx context.Context, service, proto, name string) (string, []*net.SRV, error) {
    71  	tr.mu.Lock()
    72  	defer tr.mu.Unlock()
    73  
    74  	cname := "_" + service + "._" + proto + "." + name
    75  	if srvs, ok := tr.srvLookupTable[cname]; ok {
    76  		return cname, srvs, nil
    77  	}
    78  	return "", nil, &net.DNSError{
    79  		Err:         "srvLookup error",
    80  		Name:        cname,
    81  		Server:      "fake",
    82  		IsTemporary: true,
    83  	}
    84  }
    85  
    86  func (tr *testNetResolver) LookupTXT(ctx context.Context, host string) ([]string, error) {
    87  	tr.mu.Lock()
    88  	defer tr.mu.Unlock()
    89  
    90  	if sc, ok := tr.txtLookupTable[host]; ok {
    91  		return sc, nil
    92  	}
    93  	return nil, &net.DNSError{
    94  		Err:         "txtLookup error",
    95  		Name:        host,
    96  		Server:      "fake",
    97  		IsTemporary: true,
    98  	}
    99  }
   100  
   101  func (tr *testNetResolver) UpdateTXTLookupTable(table map[string][]string) {
   102  	tr.mu.Lock()
   103  	tr.txtLookupTable = table
   104  	tr.mu.Unlock()
   105  }
   106  
   107  // txtRecordServiceConfig generates a slice of strings (aggregately representing
   108  // a single service config file) for the input config string, that represents
   109  // the result from a real DNS TXT record lookup.
   110  func txtRecordServiceConfig(cfg string) []string {
   111  	// In DNS, service config is encoded in a TXT record via the mechanism
   112  	// described in RFC-1464 using the attribute name grpc_config.
   113  	b := append([]byte("grpc_config="), []byte(cfg)...)
   114  
   115  	// Split b into multiple strings, each with a max of 255 bytes, which is
   116  	// the DNS TXT record limit.
   117  	var r []string
   118  	for i := 0; i < len(b); i += txtBytesLimit {
   119  		if i+txtBytesLimit > len(b) {
   120  			r = append(r, string(b[i:]))
   121  		} else {
   122  			r = append(r, string(b[i:i+txtBytesLimit]))
   123  		}
   124  	}
   125  	return r
   126  }
   127
View as plain text