1 /*
2  * Copyright 2018 The gRPC Authors
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package io.grpc.alts.internal;
18 
19 import com.google.common.base.Preconditions;
20 import io.netty.buffer.ByteBuf;
21 import io.netty.buffer.Unpooled;
22 import java.security.SecureRandom;
23 import java.util.ArrayList;
24 import java.util.List;
25 import java.util.Random;
26 
27 public final class ByteBufTestUtils {
28   public interface RegisterRef {
register(ByteBuf buf)29     ByteBuf register(ByteBuf buf);
30   }
31 
32   private static final Random random = new SecureRandom();
33 
34   // The {@code ref} argument can be used to register the buffer for {@code release}.
35   // TODO: allow the allocator to be passed in.
getDirectBuffer(int len, RegisterRef ref)36   public static ByteBuf getDirectBuffer(int len, RegisterRef ref) {
37     return ref.register(Unpooled.directBuffer(len));
38   }
39 
40   /** Get random bytes. */
getRandom(int len, RegisterRef ref)41   public static ByteBuf getRandom(int len, RegisterRef ref) {
42     ByteBuf buf = getDirectBuffer(len, ref);
43     byte[] bytes = new byte[len];
44     random.nextBytes(bytes);
45     buf.writeBytes(bytes);
46     return buf;
47   }
48 
49   /** Fragment byte buffer into multiple pieces. */
fragmentByteBuf(ByteBuf in, int num, RegisterRef ref)50   public static List<ByteBuf> fragmentByteBuf(ByteBuf in, int num, RegisterRef ref) {
51     ByteBuf buf = in.slice();
52     Preconditions.checkArgument(num > 0);
53     List<ByteBuf> fragmentedBufs = new ArrayList<>(num);
54     int fragmentSize = buf.readableBytes() / num;
55     while (buf.isReadable()) {
56       int readBytes = num == 0 ? buf.readableBytes() : fragmentSize;
57       ByteBuf tmpBuf = getDirectBuffer(readBytes, ref);
58       tmpBuf.writeBytes(buf, readBytes);
59       fragmentedBufs.add(tmpBuf);
60       num--;
61     }
62     return fragmentedBufs;
63   }
64 
writeSlice(ByteBuf in, int len)65   static ByteBuf writeSlice(ByteBuf in, int len) {
66     Preconditions.checkArgument(len <= in.writableBytes());
67     ByteBuf out = in.slice(in.writerIndex(), len);
68     in.writerIndex(in.writerIndex() + len);
69     return out.writerIndex(0);
70   }
71 }
72